Replace tensorboardX with torch and JAX/trax versions

mle-infrastructure / mle-toolbox

Lightweight Tool to Manage Distributed ML Experiments 🛠

https://mle-infrastructure.github.io/mle_toolbox/toolbox/

MIT License

3 stars 1 forks source link

Replace tensorboardX with torch and JAX/trax versions #31

Closed RobertTLange closed 3 years ago

RobertTLange commented 3 years ago

tensorboardX is outdated. torch provides an internal API via from torch.utils.tensorboard import SummaryWriter. Replace this and reduce dependencies. Only remaining required dependency should be tensorboard.

Furthermore, trax appears to provide a jax-compatible version as well: jaxboard with standalone installation. Check if this can simply act as a drop-in replacement with parameter FrozenDict.

RobertTLange commented 3 years ago

Addressed in 28f45f4. Example for JAX tensorboard logging can be found in examples/jax_vae. We don't log gradients since they would have to be passed seperately!