ForwardHookManager on multiple GPUs about torchdistill

[yoshitomo-matsubara]

Hi @arbellea

If you are using a module wrapped by DataParallel or DistributedDataParallel, you should provide the right device when instantiating ForwardHookManager

For instance, you can replace

device = torch.device('cpu')
forward_hook_manager = ForwardHookManager(device)

with

device = torch.device('cuda')
forward_hook_manager = ForwardHookManager(device)

in this example notebook.

[arbellea]

Thanks, but that is not what I meant
After further digging, I saw that pop_io_dict() gathers the outputs from all devices.
Since I would like to use this as a submodule of a network on parallel gpus, I need to extract only the outputs from the relevant gpu.
I wrote a small fuction that does that:

    def get_io_dict_from_device(self, device):
        device_io_dict = dict()
        for module_path, module_io_dict in self.io_dict.items():
            device_io_dict[module_path] = dict()
            for io_type in list(module_io_dict.keys()):
                sub_dict = module_io_dict[io_type]
                values = sub_dict[device.index]
                device_io_dict[module_path][io_type] = values
        return device_io_dict

and it works fine

[yoshitomo-matsubara]

@arbellea Thank you for the clarification.
Can I ask you for what purpose you'd like extract device-wise input/output?
I didn't expect the forward hook manager to be used in that way, but if that sounds useful for other users, I may want to update it.

[arbellea]

I want to take a pretrained network, let's say ResNet, and use that as a backbone for my model.
I would like to take several intermediate layers and continue computations. For example, a U-Net with a ResNet backbone.

[arbellea]

What I do now is download the code and modify it so that I can extract the intermediate layers...
But your code is much more convinent!

[yoshitomo-matsubara]

But your code is much more convinent!

Great to hear that :)
It may make sense, and I'll shortly update mine, mentioning this issue.

As you might have known, IntermediateLayerGetter in torchvision might also be something you're looking for, as it's used in R-CNN object detectors using ResNet and FPN as their backbones.