pytorch加载预训练模型

1. 查看所有module

model.modules()以深度优先遍历的方式，存储了net的所有模块，包括net itself,net's children, children of net's children。即model.children()只包括网络模块的第一代儿子模块，而model.modules()包含网络模块的自己本身和所有后代模块。

model.modules()和model.named_modules()内部采用yield关键字，得到生成器。当外部迭代调用net.named_modules()时，会先返回prefix=’’，以及net对象本身。然后下一步会递归的调用named_modules()，继而深度优先的返回每一个module。

model.modules()和model.children()只返回具体的module。这两个函数返回生成器，可以使用循环或next函数遍历:

for module in model.chilrean():
	print(module)

for module in model.modules():
	print(module)

而model.named_children()和model.named_modules()则返回名称和具体module的dict:

for name, module in model.named_children():
	print(f"{name}, {module}")

相同的操作还有named_buffer和named_parameter函数。

参考： self.modules() 和 self.children()的区别
pytorch中的named_parameters(), named_modules()

2. 删除最后一层

model = resnet50()

# method 1, ._module returns the OrderedDict
name, module = model._modules.popitem()

# method 2  
name, module = model._module.pop('fc')

# method 3
del model.fc

# method 4
model.__delattr__('fc')

但是删除最后一层之后，需要重写forward方法，因为在resnet的forward方法中有使用到fc，否则会报错has no attribute fc.

或者直接将fc换成自己的classifier:

# 添加自己的分类层
model.fc = nn.Linear()

# 换成indentity层
model.fc = IdentityLayer()

3. 使用预训练模型

有时候想要使用预训练模型的weight,但是有些自定义的部分,例如修改输入channel:

class RGBMaskEncoderCNN(nn.Module):
	def __init__(self):
		super(RGBMaskEncoderCNN, self).__init__()

		self.alexnet = models.alexnet(pretrained=True)
		# get the pre-trained weights of the first layer
		pretrained_weights = self.alexnet.features[0].weight
		new_features = nn.Sequential(*list(self.alexnet.features.children()))
		new_features[0] = nn.Conv2d(4, 64, kernel_size=11, stride=4, padding=2)

		# For M-channel weight should randomly initialized with Gaussian
		new_feaures_[0].weight.data.normal_(0, 0.001)
		# For RGB it should be copied from pretrained weights
		new_features[0].weight.data[:, :3, :, :] = pretrained_weights

		self.alexnet.features = new_features
		
	def forward(self, images):
		"""Extract Feature Vector from Input Images"""
		features = self.alexnet(images)
		return features

以上代码将输入由3改为4,但是仍然使用预训练权重。

4. 固定某些层

for param in model.parameters():
    param.requires_grad = False

5. register_forward_hook用法

在不改动网络结构的情况下获取网络中间层输出:

class Model(nn.Module):
    def __init__(self):
  	    super(Model, self).__init__()
  	    self.layer1 = Conv(...)
  	    self.layer2 = Conv(...)

	def forward(self, x):
		return self.layer2(self.layer1(x))

如果需要获取layer1的输出，直观的做法是该代码，即修改forwrd函数:

class Model(nn.Module):
    def __init__(self):
  	    super(Model, self).__init__()
  	    self.layer1 = Conv(...)
  	    self.layer2 = Conv(...)

	def forward(self, x):
		out = self.layer1(x)
		layer1_out= out
		return self.layer2(out), layer1_out

但是现在我不想改代码，因此可以使用hook从外部或者网络的中间层输出:

featuers = []
def hook(module, input, output):
	features.append(output.clone().detach())

net = Model()
x = torch.rand((2, 3, 32, 32))
handle = model.conv2.register_forward_hook(hook)
y = model(x)

print(features)
handle.remove()

得到网络的某一层，调用register_forward_hook方法，需要传入一个hook方法，该方法有三个参数:

def hook(module, input, output)

• module: 该层
• 该层输入
• 该层输出

register_forward_hook调用完成后应删除，防止每次调用增加开销。