【colab pytorch】提取模型中的某一层

时间:2022-07-23
本文章向大家介绍【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'))

随机文章
本站知识点必读
最近更新