Is there a way to log the weight changes of neural networks?


I would like to visualize how the weights are changing as a neural network continues its training. I’m prototyping this with RLlib DQN in simple Gym environments. Weights and Biases provides watch() method that logs the weight and gradients, but I’m having difficulty finding out how to integrate that into Ray Tune and RLlib.

Thanks in advance.

Hey, you could add the weights inside the on_train_result callback (override the rllib/agents/ class and implement on_train_result method). In there is a result arg, which is a dict to which you can add arbitrary data (e.g. your weights, which you can get in that method via trainer.get_policy().get_weights()).

When I tried this I was able only to get the value with shape = 1 using W&B:

There is the way I tried to logged it:

class MyCallbacks(DefaultCallbacks):
    def on_train_result(self, trainer, result: dict, **kwargs):
        for k, v in trainer.get_policy().get_weights().items():
            result["FCC/{}".format(k)] = v

How can I log also other weights other than value_out/bias?
It seems that there is skipped all other arrays.
I also tried with v.flatten() version.