题目引入
为了方便,我们用Sequential定义一个卷积神经网络中,例如
self.fc = nn.Sequential(
nn.Upsample(scale_factor=2,mode='nearest'),
nn.Conv2d(512,512,3,1,1,bias=False),
nn.BatchNorm2d(512),
nn.ReLU(True),
…………
那么我们如何查看诸如nn.Conv2d(512,512,3,1,1,bias=False)
这一层输出的大小呢?
方法就是定义一个打印函数并定义一个符合规格的假数据,然后循环输出。
代码
def printInfo(self):
x = torch.rand([3,512,1,1])
for name, module in self.fc.named_children():
x = module(x)
if isinstance(module,nn.Upsample):
print("Upsample({}) : {}".format(name,x.shape))
elif isinstance(module,nn.Conv2d):
print("Conv2d({}) : {}".format(name,x.shape))
这样,只需要使用model.printInfo()
就可以输出我们所关心的某几层的特征图规格。
如下图,第一行诸如Upsample
后面的数字就是nn.Sequential
对应的层号