pytorch加載訓練好的模型用來測試或者處理方式

更新時間：2023年09月09日 16:32:17 作者：群星閃耀

這篇文章主要介紹了pytorch加載訓練好的模型用來測試或者處理方式,具有很好的參考價值,希望對大家有所幫助,如有錯誤或未考慮完全的地方,望不吝賜教

1.直接加載預訓練模型

如果我們使用的模型和原模型完全一樣，

那么我們可以直接加載別人訓練好的模型：

import torchvision.models as models
resnet50 = models.resnet50(pretrained=True)

如果只需要網(wǎng)絡結(jié)構(gòu)，不需要用預訓練模型的參數(shù)來初始化，

那么就是：

model =torchvision.models.resnet50(pretrained=False)

2.修改某一層

PyTorch中的torchvision里已經(jīng)有很多常用的模型了，

可以直接調(diào)用：

AlexNet
VGG
ResNet
SqueezeNet
DenseNet

import torchvision.models as models
resnet18 = models.resnet18()
alexnet = models.alexnet()
squeezenet = models.squeezenet1_0()
densenet = models.densenet_161()

但是對于我們的任務而言有些層并不是直接能用，需要我們微微改一下，

比如，resnet最后的全連接層是分1000類，而我們只有21類；

又比如，resnet第一層卷積接收的通道是3，我們可能輸入圖片的通道是4，

那么可以通過以下方法修改：

resnet.conv1 = nn.Conv2d(4, 64, kernel_size=7, stride=2, padding=3, bias=False)
resnet.fc = nn.Linear(2048, 21)

3.加載部分預訓練模型

其實大多數(shù)時候我們需要根據(jù)我們的任務調(diào)節(jié)我們的模型，所以很難保證模型和公開的模型完全一樣，但是預訓練模型的參數(shù)確實有助于提高訓練的準確率，為了結(jié)合二者的優(yōu)點，就需要我們加載部分預訓練模型。

#加載model，model是自己定義好的模型
resnet50 = models.resnet50(pretrained=True) 
model =Net(...) 
#讀取參數(shù) 
pretrained_dict =resnet50.state_dict() 
model_dict = model.state_dict() 
#將pretrained_dict里不屬于model_dict的鍵剔除掉 
pretrained_dict =  {k: v for k, v in pretrained_dict.items() if k in model_dict} 
# 更新現(xiàn)有的model_dict 
model_dict.update(pretrained_dict) 
# 加載我們真正需要的state_dict
model.load_state_dict(model_dict)
# 加載我們真正需要的state_dict 
model.load_state_dict(model_dict)

4. 保存和加載自己的模型

pytorch保存模型的方式有兩種：

第一種：將整個網(wǎng)絡都都保存下來
第二種：僅保存和加載模型參數(shù)（推薦使用這樣的方法）

4.1 保存和加載整個模型

# 保存
torch.save(model_object, Path)
# 加載
model = torch.load(Path)

4.2 僅保存和加載模型參數(shù)(推薦使用)

# ----------------保存模型參數(shù)--------------------------
torch.save(model.state_dict(), PATH)
#example
torch.save(resnet50.state_dict(),'ckp/model.pth')    
# ----------------加載模型參數(shù)--------------------------
model = ModelClass(*args, **kwargs) # 這是你后來設(shè)置的模型
model.load_state_dict(torch.load(PATH)) # 加載參數(shù)
#example
resnet=resnet50(pretrained=True)
resnet.load_state_dict(torch.load('ckp/model.pth'))

4.3 每個epoch保存一個模型參數(shù)

for epoch in range(start_epoch, nEpochs + 1):
        train(training_data_loader, optimizer, model, criterion, epoch)
        save_checkpoint(model, epoch)
def save_checkpoint(model, epoch):
    model_out_path = "checkpoint/" + "model_epoch_{}.pth".format(epoch)
    state = {"epoch": epoch ,"model": model}
    if not os.path.exists("checkpoint/"):
        os.makedirs("checkpoint/")
    torch.save(state, model_out_path)
    print("Checkpoint saved to {}".format(model_out_path))

上面的代碼中start_epoch是開始保存模型的epoch,nEpochs是總共訓練的次數(shù)。

train（）里面的參數(shù)，是訓練的過程：一些訓練數(shù)據(jù)，優(yōu)化器，模型和訓練標準。