【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'))

 


免責聲明!

本站轉載的文章為個人學習借鑒使用,本站對版權不負任何法律責任。如果侵犯了您的隱私權益,請聯系本站郵箱yoyou2525@163.com刪除。



 
粵ICP備18138465號   © 2018-2025 CODEPRJ.COM