[feat] Update Aim ML adapters (#2964)

docs/example_scripts/acme_track.py

@@ -9,7 +9,7 @@
 from acme.jax import experiments
 from acme.utils import loggers
 
-from aim.sdk.acme import AimCallback, AimWriter
+from aimstack.ml.adapters.acme import AimCallback, AimWriter
 
 
 def make_environment(seed: int) -> dm_env.Environment:

docs/example_scripts/catboost_track.py

@@ -5,7 +5,7 @@
 
 import numpy as np
 
-from aim.catboost import AimLogger
+from aimstack.ml.adapters.catboost import AimLogger
 from catboost import CatBoostClassifier, Pool
 
 # initialize data

docs/example_scripts/fastai_track.py

@@ -7,7 +7,7 @@
 from fastai.vision.all import get_image_files, aug_transforms, imagenet_stats
 from fastai.vision.all import cnn_learner, resnet18, CrossEntropyLossFlat
 import regex as re
-from aim.fastai import AimCallback
+from aimstack.ml.adapters.fastai import AimCallback
 
 
 def get_arabic_mnist_labels(file_path):

docs/example_scripts/hugging_face_track.py

@@ -26,7 +26,7 @@
 
 import numpy as np
 from datasets import load_dataset, load_metric
-from aim.hugging_face import AimCallback
+from aimstack.ml.adapters.hugging_face import AimCallback
 
 import transformers
 from transformers import (

docs/example_scripts/keras_track.py

@@ -1,4 +1,4 @@
-from aim.keras import AimCallback
+from aimstack.ml.adapters.keras import AimCallback
 
 import keras
 from keras.datasets import mnist

docs/example_scripts/keras_tuner_track.py

@@ -1,6 +1,6 @@
 import tensorflow as tf
 import tensorflow_datasets as tfds
-from aim.keras_tuner import AimCallback
+from aimstack.ml.adapters.keras_tuner import AimCallback
 import kerastuner as kt
 
 
@@ -55,7 +55,7 @@ def standardize_record(record):
 tuner.search(
     train_ds,
     validation_data=test_ds,
-    callbacks=[AimCallback(repo='.', experiment_name='keras_tuner_test', tuner=tuner)],
+    callbacks=[AimCallback(experiment_name='keras_tuner_test', tuner=tuner)],
 )
 
 best_model = tuner.get_best_models(1)[0]

docs/example_scripts/lightgbm_track.py

@@ -2,7 +2,7 @@
 
 import pandas as pd
 from sklearn.metrics import mean_squared_error
-from aim.lightgbm import AimCallback
+from aimstack.ml.adapters.lightgbm import AimCallback
 
 import lightgbm as lgb
 

docs/example_scripts/mxnet_track.py

@@ -4,7 +4,7 @@
 from mxnet.gluon.model_zoo import vision
 from mxnet.gluon.contrib.estimator import estimator
 
-from aim.mxnet import AimLoggingHandler
+from aimstack.ml.adapters.mxnet import AimLoggingHandler
 
 
 gpu_count = mx.context.num_gpus()

docs/example_scripts/optuna_track.py

@@ -1,5 +1,5 @@
 import optuna
-from aim.optuna import AimCallback
+from aimstack.ml.adapters.optuna import AimCallback
 
 
 # Add Aim callback to Optuna optimization

docs/example_scripts/paddle_track.py

@@ -1,7 +1,7 @@
 import paddle
 import paddle.vision.transforms as T
 from paddle.static import InputSpec
-from aim.paddle import AimCallback
+from aimstack.ml.adapters.paddle import AimCallback
 
 
 inputs = [InputSpec([-1, 1, 28, 28], 'float32', 'image')]

docs/example_scripts/prophet_track.py

@@ -1,4 +1,4 @@
-from aim.prophet import AimLogger
+from aimstack.ml.adapters.prophet import AimLogger
 import numpy as np
 import pandas as pd
 from prophet import Prophet

docs/example_scripts/pytorch_ignite_track.py

@@ -13,7 +13,7 @@
 from ignite.handlers import global_step_from_engine, EarlyStopping
 from ignite.contrib.handlers import ProgressBar
 
-from aim.pytorch_ignite import AimLogger
+from aimstack.ml.adapters.pytorch_ignite import AimLogger
 
 # transform to normalize the data
 transform = transforms.Compose(

docs/example_scripts/pytorch_lightning_track.py

@@ -1,4 +1,4 @@
-from aim.pytorch_lightning import AimLogger
+from aimstack.ml.adapters.pytorch_lightning import AimLogger
 
 from argparse import ArgumentParser
 
@@ -47,7 +47,7 @@ def test_step(self, batch, batch_idx):
         loss = F.cross_entropy(y_hat, y)
         self.log('test_loss', loss)
         # Track metrics manually
-        self.logger.experiment.track(1, name='manually_tracked_metric')
+        self.logger.experiment.track_auto(1, name='manually_tracked_metric')
 
     def configure_optimizers(self):
         return torch.optim.Adam(self.parameters(), lr=self.hparams.learning_rate)

docs/example_scripts/pytorch_track.py

@@ -1,5 +1,5 @@
-from aim import Run
-from aim.pytorch import track_gradients_dists, track_params_dists
+from aimstack.asp import Run
+from aimstack.ml.adapters.pytorch import track_gradients_dists, track_params_dists
 
 import torch
 import torch.nn as nn
@@ -111,19 +111,19 @@ def forward(self, x):
             acc = 100 * correct / total
 
             # aim - Track metrics
-            items = {'accuracy': acc, 'loss': loss}
-            aim_run.track(items, epoch=epoch, context={'subset': 'train'})
+            aim_run.track_auto(acc, name='accuracy', epoch=epoch, context={'subset': 'train'})
+            aim_run.track_auto(loss, name='loss', epoch=epoch, context={'subset': 'train'})
 
             # aim - Track weights and gradients distributions
             track_params_dists(model, aim_run)
             track_gradients_dists(model, aim_run)
 
             # TODO: Do actual validation
             if i % 300 == 0:
-                aim_run.track(
-                    loss.item(), name='loss', epoch=epoch, context={'subset': 'val'}
+                aim_run.track_auto(
+                    loss, name='loss', epoch=epoch, context={'subset': 'val'}
                 )
-                aim_run.track(
+                aim_run.track_auto(
                     acc, name='accuracy', epoch=epoch, context={'subset': 'val'}
                 )
 

docs/example_scripts/pytorch_track_images.py

@@ -1,4 +1,4 @@
-from aim import Run
+from aimstack.asp import Run
 from aimstack.asp.models.objects.image import convert_to_aim_image_list
 
 import torch
@@ -105,7 +105,7 @@ def forward(self, x):
             )
 
             # aim - Track model loss function
-            aim_run.track(
+            aim_run.track_auto(
                 loss.item(), name='loss', epoch=epoch, context={'subset': 'train'}
             )
 
@@ -117,23 +117,23 @@ def forward(self, x):
             acc = 100 * correct / total
 
             # aim - Track metrics
-            aim_run.track(
+            aim_run.track_auto(
                 acc, name='accuracy', epoch=epoch, context={'subset': 'train'}
             )
 
-            aim_run.track(
+            aim_run.track_auto(
                 aim_images, name='images', epoch=epoch, context={'subset': 'train'}
             )
 
             # TODO: Do actual validation
             if i % 300 == 0:
-                aim_run.track(
-                    loss.item(), name='loss', epoch=epoch, context={'subset': 'val'}
+                aim_run.track_auto(
+                    loss, name='loss', epoch=epoch, context={'subset': 'val'}
                 )
-                aim_run.track(
+                aim_run.track_auto(
                     acc, name='accuracy', epoch=epoch, context={'subset': 'val'}
                 )
-                aim_run.track(
+                aim_run.track_auto(
                     aim_images, name='images', epoch=epoch, context={'subset': 'val'}
                 )
 

docs/example_scripts/sb3_track.py

@@ -1,4 +1,4 @@
-from aim.sb3 import AimCallback
+from aimstack.ml.adapters.sb3 import AimCallback
 from stable_baselines3 import A2C
 
 

docs/example_scripts/tensorflow_1_keras_track.py

@@ -1,4 +1,4 @@
-from aim.tensorflow import AimCallback
+from aimstack.ml.adapters.tensorflow import AimCallback
 
 import tensorflow as tf
 

docs/example_scripts/tensorflow_keras_track.py

@@ -1,4 +1,4 @@
-from aim.tensorflow import AimCallback
+from aimstack.ml.adapters.tensorflow import AimCallback
 
 import tensorflow as tf
 

docs/example_scripts/xgboost_track.py

@@ -5,7 +5,7 @@
 
 import numpy as np
 import xgboost as xgb
-from aim.xgboost import AimCallback
+from aimstack.ml.adapters.xgboost import AimCallback
 
 # label need to be 0 to num_class -1
 data = np.loadtxt(

pkgs/aimstack/asp/loggers/metric.py

@@ -1,6 +1,8 @@
 import numbers
 import json
 
+from typing import Optional, Any
+
 from typing import TYPE_CHECKING
 if TYPE_CHECKING:
     from pandas import DataFrame
@@ -64,7 +66,11 @@ def dataframe(
 
 
 class Metric(Sequence[numbers.Number], DataframeMixin):
-    ...
+    def track(self, value: Any, *, step: Optional[int] = None, **axis):
+        from aim._sdk import num_utils
+        if num_utils.is_number(value):
+            value = num_utils.convert_to_py_number(value)
+        super().track(value, step=step, **axis)
 
 
 class SystemMetric(Sequence[numbers.Number], DataframeMixin):

pkgs/aimstack/asp/loggers/run.py

@@ -6,13 +6,14 @@
 
 from functools import partialmethod
 
-from aim import Container, Property
+from aim import Container, Sequence, Property
 from aim._sdk.utils import utc_timestamp
 from aim._sdk import type_utils
 from aimcore.callbacks import Caller
 from aimcore.callbacks import events
 from aim._ext.system_info import utils as system_utils
 from aim._sdk.constants import ContainerOpenMode, KeyNames
+from aim._sdk.num_utils import is_number
 
 from .logging import (
     LogLine,
@@ -28,7 +29,7 @@
 from .distribution import DistributionSequence
 from .figures import FigureSequence, Figure3DSequence
 
-from typing import Optional, Union, List, Tuple, Dict, Any
+from typing import Optional, Union, List, Tuple, Dict, Any, Type
 
 from typing import TYPE_CHECKING
 if TYPE_CHECKING:
@@ -135,6 +136,12 @@ def track(self, value, name: str, step: Optional[int] = None, context: dict = Op
         sequence = self.sequences[name, context]
         sequence.track(value, step=step, **axis)
 
+    def track_auto(self, value, name: str, step: int = None, context: dict = None, **axis):
+        context = {} if context is None else context
+        seq_type = self._get_sequence_type_from_value(value)
+        sequence = self.sequences.typed_sequence(seq_type, name, context)
+        sequence.track(value, step=step, **axis)
+
     def get_metric(self, name: str, context: Optional[dict] = None) -> Metric:
         context = {} if context is None else context
         return self.sequences.typed_sequence(Metric, name, context)
@@ -204,3 +211,23 @@ def dataframe(
         import pandas as pd
         df = pd.DataFrame(data, index=[0])
         return df
+
+    @staticmethod
+    def _get_sequence_type_from_value(value) -> Type[Sequence]:
+        val_type = type_utils.get_object_typename(value)
+        if type_utils.is_allowed_type(val_type, type_utils.get_sequence_value_types(Metric)) \
+                or is_number(value):
+            return Metric
+        if type_utils.is_allowed_type(val_type, type_utils.get_sequence_value_types(ImageSequence)):
+            return ImageSequence
+        if type_utils.is_allowed_type(val_type, type_utils.get_sequence_value_types(AudioSequence)):
+            return AudioSequence
+        if type_utils.is_allowed_type(val_type, type_utils.get_sequence_value_types(TextSequence)):
+            return TextSequence
+        if type_utils.is_allowed_type(val_type, type_utils.get_sequence_value_types(DistributionSequence)):
+            return TextSequence
+        if type_utils.is_allowed_type(val_type, type_utils.get_sequence_value_types(FigureSequence)):
+            return TextSequence
+        if type_utils.is_allowed_type(val_type, type_utils.get_sequence_value_types(Figure3DSequence)):
+            return TextSequence
+        return Sequence

pkgs/aimstack/ml/__init__.py

@@ -0,0 +1,6 @@
+from .loggers.run import Run
+
+
+__all__ = ['Run']
+
+__aim_types__ = [Run]