Pytorch Forward Hook

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.

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.

1 Like

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