I have followed this forward-and-backward-function-hooks to print the layer’s feature map’s size,just like:
Conv2d forward input: torch.Size([32, 3, 300, 300]) output: torch.Size([32, 64, 300, 300])
Conv2d forward input: torch.Size([32, 64, 300, 300]) output: torch.Size([32, 64, 300, 300])
but this “Conv2d” name is unrecognized and I want to print the name of the layer, how can i approach this?
for example:
when i define a module, i use:
def conv_layers():
conv1 = nn.Conv2d(1, 10, 5)
pool1 = nn.MaxPool2d(2, 2)
conv2 = nn.Conv2d(10, 20, 5)
pool2 = nn.MaxPool2d(2, 2)
layers = [conv1, pool1, conv2, pool2]
return layers
class LeNet(nn.Module):
def __init__(self,base):
self.base= nn.Sequential(*base)
self.fc1 = nn.Linear(320, 50)
self.fc2 = nn.Linear(50, 10)
def forward():
……
net=LeNet(conv_layers)
def printnorm(self, input, output):
print('{} forward\t input: {}\t output: {}\t'.format(
self.__class__.__name__, input[0].size(), output.data.size()))
layers = []
for layer in net.children():
layer.register_forward_hook(printnorm)
out = net(input)
I get something like this:
Conv2d forward input: torch.Size([32, 3, 300, 300]) output: torch.Size([32, 64, 300, 300])
ReLU forward input: torch.Size([32, 64, 300, 300]) output: torch.Size([32, 64, 300, 300])
but how can i know the Conv2d’s name？
because the layers in “self.base” are:
base[0]->conv2d ,base[1]->maxpool2d, base[2]->conv2d ,base[3]->maxpool2d,
I want get the full name in the hook function “printnorm” so i can find the specific Conv2d and do something else.