wehub-resource-sync
94 lines
3.3 KiB
Python
94 lines
3.3 KiB
Python
"""Integration with [Comet.ml](https://www.comet.ml/).
|
|
|
|
Docs: https://docs.fast.ai/callback.comet.html.md"""
|
|
|
|
# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/70d_callback.comet.ipynb.
|
|
|
|
# %% auto #0
|
|
__all__ = ['CometCallback']
|
|
|
|
# %% ../../nbs/70d_callback.comet.ipynb #d99c2bf2
|
|
import tempfile
|
|
|
|
from ..basics import *
|
|
from ..learner import Callback
|
|
|
|
# %% ../../nbs/70d_callback.comet.ipynb #f9458afe
|
|
import comet_ml
|
|
|
|
# %% ../../nbs/70d_callback.comet.ipynb #81b59bd6
|
|
class CometCallback(Callback):
|
|
"Log losses, metrics, model weights, model architecture summary to neptune"
|
|
order = Recorder.order + 1
|
|
|
|
def __init__(self, project_name, log_model_weights=True):
|
|
self.log_model_weights = log_model_weights
|
|
self.keep_experiment_running = keep_experiment_running
|
|
self.project_name = project_name
|
|
self.experiment = None
|
|
|
|
def before_fit(self):
|
|
try:
|
|
self.experiment = comet_ml.Experiment(project_name=self.project_name)
|
|
except ValueError:
|
|
print("No active experiment")
|
|
|
|
try:
|
|
self.experiment.log_parameter("n_epoch", str(self.learn.n_epoch))
|
|
self.experiment.log_parameter("model_class", str(type(self.learn.model)))
|
|
except:
|
|
print(f"Did not log all properties.")
|
|
|
|
try:
|
|
with tempfile.NamedTemporaryFile(mode="w") as f:
|
|
with open(f.name, "w") as g:
|
|
g.write(repr(self.learn.model))
|
|
self.experiment.log_asset(f.name, "model_summary.txt")
|
|
except:
|
|
print("Did not log model summary. Check if your model is PyTorch model.")
|
|
|
|
if self.log_model_weights and not hasattr(self.learn, "save_model"):
|
|
print(
|
|
"Unable to log model to Comet.\n",
|
|
)
|
|
|
|
def after_batch(self):
|
|
# log loss and opt.hypers
|
|
if self.learn.training:
|
|
self.experiment.log_metric("batch__smooth_loss", self.learn.smooth_loss)
|
|
self.experiment.log_metric("batch__loss", self.learn.loss)
|
|
self.experiment.log_metric("batch__train_iter", self.learn.train_iter)
|
|
for i, h in enumerate(self.learn.opt.hypers):
|
|
for k, v in h.items():
|
|
self.experiment.log_metric(f"batch__opt.hypers.{k}", v)
|
|
|
|
def after_epoch(self):
|
|
# log metrics
|
|
for n, v in zip(self.learn.recorder.metric_names, self.learn.recorder.log):
|
|
if n not in ["epoch", "time"]:
|
|
self.experiment.log_metric(f"epoch__{n}", v)
|
|
if n == "time":
|
|
self.experiment.log_text(f"epoch__{n}", str(v))
|
|
|
|
# log model weights
|
|
if self.log_model_weights and hasattr(self.learn, "save_model"):
|
|
if self.learn.save_model.every_epoch:
|
|
_file = join_path_file(
|
|
f"{self.learn.save_model.fname}_{self.learn.save_model.epoch}",
|
|
self.learn.path / self.learn.model_dir,
|
|
ext=".pth",
|
|
)
|
|
else:
|
|
_file = join_path_file(
|
|
self.learn.save_model.fname,
|
|
self.learn.path / self.learn.model_dir,
|
|
ext=".pth",
|
|
)
|
|
self.experiment.log_asset(_file)
|
|
|
|
def after_fit(self):
|
|
try:
|
|
self.experiment.end()
|
|
except:
|
|
print("No neptune experiment to stop.")
|