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
调用完成后应删除,防止每次调用增加开销。