Enable training with Tensorboard tracking #209
Conversation
Co-authored-by: Andrew <daia99@users.noreply.github.com>
|
We may want to remove the tensorboard files for the PR! |
woops, got it |
| @@ -78,19 +78,32 @@ def __init__(self, config, **kwargs): | |||
| dist_config = get_distributed_config(self.accelerator) | |||
| config_dict["distributed"] = dist_config | |||
| init_trackers_kwargs = {} | |||
| if "wandb" in config.train.trackers: | |||
| # HACK: Tensorboard doesn't like nested dict as hyperparams | |||
| config_dict_flat = {a:b for (k,v) in config_dict.items() for (a,b) in v.items() if not isinstance(b, dict)} | |||
There was a problem hiding this comment.
You can use trlx.utils.modeling.flatten_dict here
There was a problem hiding this comment.
thanks for the suggestion! I replaced the dict comprehension with a call to flatten_dict(), since tensorboard also doesn't like lists I added a couple lines to split the optimizer betas
| config=config_dict, | ||
| init_kwargs=init_trackers_kwargs, | ||
| ) | ||
| else: |
There was a problem hiding this comment.
Thanks!
Could you add back a comment explaining what this branch is for and the flattening? Aside from that it LGTM
There was a problem hiding this comment.
Sure!
- I had run into problems running trlx without wandb (I don't have an account as of now), and found an already opened issue precisely on this. The branch has minor modifications (which don't change the previous interface for wandb users) to allow tensorboard tracking
- The only tricky part is that wandb is pretty fancy and take nested dicts as logging params, this is not the case for tensorboard, hence the experiment config is fully flattened, and the only list is simply split apart (for the same reason)
Do let me know if anything is not clear or if I should add comments in the tensorboard specific logging
There was a problem hiding this comment.
Oh sorry, by branch I meant the else branch - i.e a short comment like
else: # tracker == 'tensorboard' and # flatten config for tensorboard, flatten lists in hparams into flattened config
There was a problem hiding this comment.
Got it, I was reading the comment in between my regular work and got fully confused, thanks for the feedback
|
Thanks for contributing! This LGTM now. |
This bug was introduced by CarperAI#209, which changed the `trackers` config property to `tracker` but didn't update this use case.
This bug was introduced by #209, which changed the `trackers` config property to `tracker` but didn't update this use case.
Currently only wandb logging is supported, and using tensorboard will result in a number of small blockers.
This small PR removes the following blockers:
init_trackersto check the specified trackerDisclaimers
CC @daia99