Useful code to save pytorch training to tensorboard

tensorboard_torch_histogram.py

from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter()

def save_histograms(writer, step, model):
    def flatten_gru_weights(gru_layer):
        flattened_weights = []
        for param in gru_layer.parameters():
            flattened_weights.append(param.data.view(-1))
        return torch.cat(flattened_weights)
    
    def save_layer(layer, name, step):
        weights = layer.weight
        flattened_weights = weights.flatten()
        tag = f"layer_{name}/weights"
        writer.add_histogram(tag, flattened_weights, global_step=step, bins='tensorflow')
        bias = layer.bias
        flattened_bias = bias.flatten()
        tag = f"layer_{name}/bias"
        writer.add_histogram(tag, flattened_bias, global_step=step, bins='tensorflow')
    
    def save_layers(layers, name, step):
        for index, layer in enumerate(layers):
            if isinstance(layer, nn.Linear):
                save_layer(layer, name + '/' + str(index), step)
    save_layers(model.layers)