Pytorch Forward Hook

[Juz_Z]

I am trying to extract activation vectors from some spotlighted layers using hook.
Referring to the usage from the forum, I applied this hook function to my resnet50 model (trained with dataparallel, multi-GPUs).

While I am sending batched test images to the network and expecting to get batched feature vectors, I am surprised that the size of the activations is not the same as my batch size. I use 64 images and the output logits form the network is also 64. But the activation tensor is only 32.

I don’t know what made this failure. Please help me out.

[Diego]

You are using nn.DataParallel so the batch is split evenly across all GPUs so the initial 64 batch size is split between both gpus into 32 and 32. I think that this might be the problem in your case.

[Juz_Z]

Exactly. After observing the outputs, I realized that I used two gpus in parallel. Thank you so much for quick reply.