From 2d7ad8119c5962a320e6ccc39407e417d088c525 Mon Sep 17 00:00:00 2001 From: Piotr Kozakowski Date: Fri, 3 Nov 2017 19:50:11 +0100 Subject: Add CometML integration --- trainer/plugins.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) (limited to 'trainer') diff --git a/trainer/plugins.py b/trainer/plugins.py index 552804a..f8c299b 100644 --- a/trainer/plugins.py +++ b/trainer/plugins.py @@ -265,3 +265,23 @@ class StatsPlugin(Plugin): @classmethod def _fields_to_pairs(cls, fields): return list(map(cls._field_to_pair, fields)) + + +class CometPlugin(Plugin): + + def __init__(self, experiment, fields): + super().__init__([(1, 'epoch')]) + + self.experiment = experiment + self.fields = [ + field if type(field) is tuple else (field, 'last') + for field in fields + ] + + def register(self, trainer): + self.trainer = trainer + + def epoch(self, epoch_index): + for (field, stat) in self.fields: + self.experiment.log_metric(field, self.trainer.stats[field][stat]) + self.experiment.log_epoch_end(epoch_index) -- cgit v1.2.3-70-g09d2