【colab pytorch】提取模型中的某一层
modules()会返回模型中所有模块的迭代器,它能够访问到最内层,比如self.layer1.conv1这个模块,还有一个与它们相对应的是name_children()属性以及named_modules(),这两个不仅会返回模块的迭代器,还会返回网络层的名字。
# 取模型中的前两层 new_model = nn.Sequential(*list(model.children())[:2] # 如果希望提取出模型中的所有卷积层,可以像下面这样操作: for layer in model.named_modules(): if isinstance(layer[1],nn.Conv2d): conv_model.add_module(layer[0],layer[1])
部分层使用预训练模型:
model.load_state_dict(torch.load('model.pth'), strict=False)
注意如果保存的模型是 torch.nn.DataParallel,则当前的模型也需要是。
将GPU保存的模型加载到CPU:
model.load_state_dict(torch.load('model.pth', map_location='cpu'))