PyTorch 加載預訓練權重

Question

我正在嘗試將預訓練的 model re.net_18.pth文件加載到 pytorch 中。在線文檔建議這樣導入：

weights = torch.load("resnet_18.pth")

當我打印weights的 output 時，它給出如下內容：

 ('module.layer4.1.bn2.running_mean', tensor([ 9.1797e+01, -2.4204e+02,  5.6480e+01, -2.0762e+02,  4.5270e+01,
        -3.2356e+02,  1.8662e+02, -1.4498e+02, -2.3701e+02,  3.2354e+01,
...

所有教程都提到使用 model 為基礎加載權重：

model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.eval()

我想使用默認的 re.net-18 model 來應用權重，但是我來自tensorflow vision的resent18沒有load_state_dict function。感謝幫助。

from torchvision.models import resnet18
resnet18.load_state_dict(torch.load("resnet_18.pth"))

# 'function' object has no attribute 'load_state_dict'

Answer 1

re.net18本身是一個 function，它返回一個ResNet18 model。加載自己的預訓練權重可以做的是使用

model = resnet18()
model.load_state_dict(torch.load("resnet_18.pth"))

請注意， load_state_dict(...)就地加載權重並且不會返回 model 本身。