我們?cè)趯W(xué)習(xí)pytorch時(shí)，不可避免的要加載不同的預(yù)訓(xùn)練模型。而且pytorch下的預(yù)訓(xùn)練模型有很多種形式，我們又該如何加載呢。今天，我就為大家介紹三種常用的模型形式以及其加載方式。

1.pth形式和.pt形式，這種形式的模型使我們較為常見的形式，保存的可以是整個(gè)模型（含模型結(jié)構(gòu)及參數(shù)）或只保存參數(shù)。
當(dāng)保存的是整個(gè)模型時(shí)，讀取方式有兩種：
? ?1）.直接加載：

if __name__=='__main__':with torch.no_grad():model=InceptionI3d(num_classes=400,in_channels = 3)if not os.path.exists('./rgb_imagenet.pt'):print ('No weights Found! please download first, or comment 382~384th line')model.load_state_dict('./rgb_imagenet.pt')#通過load_state_dict()函數(shù)來加載

? ?2）.通過state_dict()函數(shù)先讀參數(shù)再保存

if __name__=='__main__':with torch.no_grad():model=InceptionI3d(num_classes=400,in_channels = 3)if not os.path.exists('./rgb_imagenet.pt'):print ('No weights Found! please download first, or comment 382~384th line')weight = torch.load("./rgb_imagenet.pt")['model']# 'model'為字典鍵值model.load_state_dict(weight.state_dict())

當(dāng)僅保存模型參數(shù)時(shí)：?這種形式一般保存的是一個(gè)字典型state_dict:權(quán)重，所以要加工一次。

if __name__ == "__main__":with torch.no_grad():net = TSN(num_class, this_test_segments if is_shift else 1, modality,base_model='resnet50',consensus_type = crop_fusion_type,img_feature_dim = img_feature_dim,pretrain = pretrain,is_shift=is_shift, shift_div=shift_div, shift_place=shift_place,non_local='_nl' in this_weights,)weights = 'TSM_something_RGB_resnet50_shift8_blockres_avg_segment16_e45.pth'#指定路徑checkpoint = torch.load(this_weights)#通過load函數(shù)讀出來checkpoint = checkpoint['state_dict']#取出state_sict所對(duì)應(yīng)的權(quán)重名稱和權(quán)重，這里有可能字典鍵值不是state_dict，可以使用.keys()函數(shù)查看具體鍵值base_dict = {'.'.join(k.split('.')[1:]): v for k, v in list(checkpoint.items())}#一般是為了除去權(quán)重名稱前的module.前綴，這個(gè)可根據(jù)自己的需要添加。也可以打印出所有的權(quán)重名稱.items()。net.load_state_dict(base_dict)#同樣用load_state_dict加載。

2.pth.tar形式，這種形式的模型保存的通常也是字典型，但是不僅僅state_sict一項(xiàng)，可以用print(set(checkpoint))打印查看，我們可以不管其他內(nèi)容，因此和2的處理方式基本相同。

if __name__ == "__main__":with torch.no_grad():model = MultiColumn(174, Model, 512)checkpoint_path = './model_best.pth.tar'#指定路徑checkpoint = torch.load(checkpoint_path)#使用load函數(shù)讀值#print(set(checkpoint))checkpoint = checkpoint['state_dict']base_dict = {'.'.join(k.split('.')[1:]): v for k, v in list(checkpoint.items())}#一般是為了除去權(quán)重名稱前的module.前綴，這個(gè)可根據(jù)自己的需要添加。可以打印出所有的權(quán)重稱.items()。model.load_state_dict(checkpoint['state_dict'])#使用load_state_dict()加載權(quán)重。

這里的介紹就先到這里啦，大家要學(xué)會(huì)靈活運(yùn)用。

總結(jié)

以上是生活随笔為你收集整理的pytorch 之 加载不同形式的预训练模型的全部?jī)?nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。