From 2ceb54fbc464384fd63b429ade75357182aa479d Mon Sep 17 00:00:00 2001 From: Dushyant Behl Date: Tue, 12 Mar 2024 14:41:29 +0530 Subject: [PATCH] custom aim callback should be minimal Signed-off-by: Dushyant Behl --- tuning/trackers/aimstack_tracker.py | 26 ++++++++++---------------- 1 file changed, 10 insertions(+), 16 deletions(-) diff --git a/tuning/trackers/aimstack_tracker.py b/tuning/trackers/aimstack_tracker.py index 30c000cc5..82f424f36 100644 --- a/tuning/trackers/aimstack_tracker.py +++ b/tuning/trackers/aimstack_tracker.py @@ -40,18 +40,6 @@ def on_train_begin(self, args, state, control, model=None, **kwargs): # call directly to make sure hyper parameters and model info is recorded. self.setup(args=args, state=state, model=model) - def track_metrics(self, metric, name, context): - run = self.experiment - if run is not None: - run.track(metric, name=name, context=context) - - - def set_params(self, params, name): - run = self.experiment - if run is not None: - [run.set((name, key), value, strict=False) for key, value in params.items()] - - class AimStackTracker(Tracker): def __init__(self, tracker_config: AimConfig): super().__init__(name="aim", tracker_config=tracker_config) @@ -79,10 +67,16 @@ def get_hf_callback(self): def track(self, metric, name, stage="additional_metrics"): context = {"subset": stage} - self.hf_callback.track_metrics(metric, name=name, context=context) + callback = self.hf_callback + run = callback.experiment + if run is not None: + run.track(metric, name=name, context=context) def set_params(self, params, name="extra_params"): try: - self.hf_callback.set_params(params, name) - except: - pass + callback = self.hf_callback + run = callback.experiment + if run is not None: + [run.set((name, key), value, strict=False) for key, value in params.items()] + except Exception as e: + raise e