Is there a way to plot the histogram of the gradients ?

[TrayfourYu]

I try to find a mount point for hooks to plot the histogram of the gradients . Just similiar to the question.

There is not mount point between loss.backward() and optimizer.step since they are wrappered deeper than the trainloop.

And, the way proposed by @zhouzaida

The least intrusive way is to inherit from MMDistributedDataParallel and override the train_step method

likely can't get the Visualizer easily. So, is there a way to do this? Thanks in advance.

[TrayfourYu]

I think i figure it out. Just customlize a new DistributedDataParallell class override from MMDistributedDataParallel and Visualizer is available though get_current_instance(). Then use the add_histogram of the TensorboardVisBackend.

from typing import Any, Dict, Union
import torch
from mmengine.optim import OptimWrapper
from mmengine.registry import MODEL_WRAPPERS
from mmengine.model.utils import detect_anomalous_params
from mmengine.visualization import Visualizer
from mmengine.logging import MMLogger
from mmengine.model.wrappers.distributed import MMDistributedDataParallel
import numpy as np

@MODEL_WRAPPERS.register_module()
class CustomDistributedDataParallel(MMDistributedDataParallel):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.step_count = 0
    
    def train_step(self, data: Union[dict, tuple, list],
                   optim_wrapper: OptimWrapper) -> Dict[str, torch.Tensor]:
        # Enable automatic mixed precision training context.
        with optim_wrapper.optim_context(self):
            data = self.module.data_preprocessor(data, training=True)
            losses = self._run_forward(data, mode='loss')
        parsed_loss, log_vars = self.module.parse_losses(losses)

        loss = optim_wrapper.scale_loss(parsed_loss)
        optim_wrapper.backward(loss)

        # modify gradients here
        model = self.module
        visualizer_now = Visualizer.get_current_instance()
        logger = MMLogger.get_current_instance()
        if self.step_count % 100 == 0 and (torch.distributed.get_rank() == 0):
            for name, param in model.named_parameters():
                if param.requires_grad and param.grad is not None:
                    grad = param.grad.detach().cpu().numpy().flatten()
                    
                    try:
                        visualizer_now.get_backend('TensorboardVisBackend').experiment.add_histogram(
                            tag=f'grad_hist/{name}',
                            values=grad,
                            global_step=self.step_count
                        )
                    except Exception as e:
                        logger.error(f"Adding histograms failed: {name} -> {str(e)}")

        if optim_wrapper.should_update():
            optim_wrapper.step()
            optim_wrapper.zero_grad()
            self.step_count += 1

        if torch.distributed.is_initialized():
            dist_tensor = torch.tensor([self.step_count], device='cuda')
            torch.distributed.broadcast(dist_tensor, src=0)
            self.step_count = dist_tensor.item()

        if self.detect_anomalous_params:
            detect_anomalous_params(parsed_loss, model=self)
        return log_vars

Config file alse need modifyed:

custom_imports = dict(
    imports=['PATH.TO.CustomDistributedDataParallel'],
    allow_failed_imports=False)
model_wrapper_cfg = dict(
    type='CustomDistributedDataParallel',
)

Hear we have Train code for one node and multi GPU
bash tools/dist_train.sh $PATH/TO/config.py $NUMOFGPUS

Finally, use Tensorboard by tensorboard --logdir $PATH/TO/FILE --port $YOURPORT and open http://localhost:$YOURPORT in your browser.