Remove tensorflow dependency for `TensorboardTracker`

[vwxyzjn]

Proposed refactoring or deprecation

Cool project! I was wondering if it's possible to remove the tensorflow dependency when using TensorboardTracker.

https://github.com/aimhubio/aim/blob/3aee452d1a68c0ff77b75ba7a0dd19dcdd095042/pkgs/aimstack/tensorboard_sync/tracker.py#L3

Motivation

tensorflow is a heavy dependency that we would like to avoid pulling if there is no real need for it.

Pitch

You can probably import tensorflow as tf only when it's needed here:

https://github.com/aimhubio/aim/blob/3aee452d1a68c0ff77b75ba7a0dd19dcdd095042/pkgs/aimstack/tensorboard_sync/tracker.py#L184-L196

[mihran113]

Hey @vwxyzjn! Thanks a lot for the report. Actually tensorboard is not listed as a dependency for aim package and it's only imported in tensorboard_sync/tracker.py and that is only used for real-time conversion of tensorboard logs so won't be a part of general aim flow and won't cause any ImportErrors if tensorboard is not installed.

[vwxyzjn]

Hi @mihran113 thanks for the reply. I meant tensorflow instead of tensorboard.

[mihran113]

@vwxyzjn sorry for misreading, but same goes for tensorflow as well, we don't have it as a dependency. When installing aim tensorflow won't be installed as a dependency and it won't cause any issues unless you're using tensorboard_sync.run.Run objects.

[vwxyzjn]

@mihran113 that is correct. However, we only want to do the following without using the tensorflow dependency.

from aim.ext.tensorboard_tracker import Run as AimRun
aim_run = AimRun(sync_tensorboard_log_dir=f"runs/{run_name}")
writer.add_scalar("charts/learning_rate", optimizer.param_groups[0]["lr"], global_step)

For additional context, we at CleanRL have been using tensorboard to log metrics, with an option to toggle tensorboard syncing to wandb. Some of our users have been asking integration with aim, so we thought of the following solution that optionally supports both aim and wandb.

    run_name = f"{args.env_id}__{args.exp_name}__{args.seed}__{int(time.time())}"
    if args.wandb:
        import wandb
        wandb.init(
            project=args.wandb_project_name,
            entity=args.wandb_entity,
            sync_tensorboard=True,
            config=vars(args),
            name=run_name,
            monitor_gym=True,
            save_code=True,
        )
    writer = SummaryWriter(f"runs/{run_name}")
    if args.aim:
        from aim.ext.tensorboard_tracker import Run as AimRun
        aim_run = AimRun(sync_tensorboard_log_dir=f"runs/{run_name}")

That said, I noticed from aim.ext.tensorboard_tracker import Run would require tensorflow dependency, which we previously did not require for wandb's tensorboard sync. This is kind of a niche request, but I was wondering if it's possible to remove the tensorflow dependency.

[mihran113]

@vwxyzjn gotcha! Makes sense, will try to squeeze it in the next patch release 🙌

[vwxyzjn]

Thanks!