WeightsHistHandler should plot all weights (incl those without grad)

[vfdev-5]

🚀 Feature

Currently, WeightsHistHandler do not log weights without grad:

ignite/ignite/contrib/handlers/tensorboard_logger.py

Lines 434 to 435 in a168476

	if p.grad is None:
	continue

Let's put an option to enable logging all weights. Maybe, we can make it even without option and just log everything ?

Meanwhile, here is a workaround code for TensorboardLogger:

from ignite.contrib.handlers import TensorboardLogger
from ignite.contrib.handlers.tensorboard_logger import WeightsHistHandler

class FixedWeightsHistHandler(WeightsHistHandler):
    
    def __call__(self, engine, logger, event_name):
        if not isinstance(logger, TensorboardLogger):
            raise RuntimeError("Handler 'WeightsHistHandler' works only with TensorboardLogger")

        global_step = engine.state.get_event_attrib_value(event_name)
        tag_prefix = f"{self.tag}/" if self.tag else ""
        for name, p in self.model.named_parameters():            
            name = name.replace(".", "/")
            logger.writer.add_histogram(
                tag=f"{tag_prefix}weights/{name}", values=p.data.detach().cpu().numpy(), global_step=global_step,
            )

[DhDeepLIT]

Assume you have a partial frozen model, it makes sense to only trace the evolution of the weight which are updated.
The question is does checking if p.grad is not None is what we want ?
What about weights modification with other computation ? like pca_low_rank ?

[vfdev-5]

I see what you mean. Ultimately, we would like to have an optional list of weights to plot ?
Something like:

from ignite.contrib.handlers import TensorboardLogger
from ignite.contrib.handlers.tensorboard_logger import WeightsHistHandler


class FixedWeightsHistHandler(WeightsHistHandler):

    def __init__(self, model, tag=None, names_whitelist=None):
        super(WeightsHistHandler, self).__init__(model, tag=tag)
        self.names_whitelist = names_whitelist
    
    def __call__(self, engine, logger, event_name):
        if not isinstance(logger, TensorboardLogger):
            raise RuntimeError("Handler 'WeightsHistHandler' works only with TensorboardLogger")

        global_step = engine.state.get_event_attrib_value(event_name)
        tag_prefix = f"{self.tag}/" if self.tag else ""
        for name, p in self.model.named_parameters():
            
            if self.names_whitelist is not None and name not in self.names_whitelist:
                continue

            name = name.replace(".", "/")
            logger.writer.add_histogram(
                tag=f"{tag_prefix}weights/{name}", values=p.data.detach().cpu().numpy(), global_step=global_step,
            )


tb_logger = TensorboardLogger(log_dir="experiments/tb_logs")
# ...
tb_logger.attach(
    trainer,
    event_name=Events.ITERATION_COMPLETED,
    log_handler=FixedWeightsHistHandler(model, names_whitelist=[n for n, _ in model.named_parameters() if "conv" in n])
)

[vfdev-5]

@DhDeepLIT I thought you wanted to avoid the case of inexisting grads and log weights in any case ?

if p.grad is None:
    continue

Or you are now talking about GradsHistHandler ?

[DhDeepLIT]

Yes that is my need, but on your side, what was THE thing you wanted to do ?
Consider that the zero_grad() function can also set grad to None.
Thus checking if the grad attribute is None or Not may not be the right option.

[DhDeepLIT]

BTW your proposed patch is fine for my need, the whitelist is a great idea too.

[vfdev-5]

FYI: zero_grad does not set grads to None by default but just zeros them :)

https://github.com/pytorch/pytorch/blob/33e9a0b5f6674bd429bda689534e1c987b38cf6e/torch/nn/modules/module.py#L1785-L1794

[DhDeepLIT]

Maybe it should be the other way around though :
it is easier to filter the weights/grad we want if a key is present in the module full name say you want to check ALL fc layer. It would be better to check if 'fc' is in the named parameter of the module, right ?

[vfdev-5]

I agree that non-exact match for names_whitelist could be helpful...

[DhDeepLIT]

https://pytorch.org/docs/stable/generated/torch.optim.Optimizer.zero_grad.html

There is a option with 'set_to_none=True'

[sadra-barikbin]

Hi @vfdev-5
I'm back. Assign it to me if you will.