使用预训练模型

使用torch.hub.load可以一行代码直接在github上下载和调用预训练好的模型到本地中。

import torch
model = torch.hub.load('pytorch/vision', 'alexnet', pretrained=True)
model.eval()

也可以自己将github上预训练好的权重文件下载到本地，再操作它。

这个链接提供了一些预训练模型下载地址。

https://blog.csdn.net/weixin_42535423/article/details/108449854#comments_19068404

但是自己使用resnext101_32x8d时，因为最近github很难打开，在自己服务器上根本ping不动。更别说400M的文件了。

所以可以自己用科学上学的办法，最好是买个Ghelper的VIP，下载大文件很快。

我自己用免费版本不能下大文件。

torch.hub.load可以从本地加载，请注意看官方文档：

https://pytorch.org/docs/stable/hub.html

下载后的注意是权重文件，所以要在自己本地先引入模型，再传入参数。

import torch
from torchvision.models.resnet import *
checkpoint = '../checkpoints/resnext101_32x8d_wsl_checkpoint.pth'
state_dict = torch.load(checkpoint)
model = resnext101_32x8d(pretrained=False)
model.load_state_dict(state_dict)

model.fc = torch.nn.Linear(2048, n_classes)
model.to(device)

主要看上面这段。

其中model = resnext101_32x8d(pretrained=False)是支持哪些模型的，可以看：

https://github.com/pytorch/vision/blob/main/hubconf.py

这里贴出来：

# Optional list of dependencies required by the package
dependencies = ["torch"]

# classification
from torchvision.models.alexnet import alexnet
from torchvision.models.densenet import densenet121, densenet169, densenet201, densenet161
from torchvision.models.efficientnet import (
    efficientnet_b0,
    efficientnet_b1,
    efficientnet_b2,
    efficientnet_b3,
    efficientnet_b4,
    efficientnet_b5,
    efficientnet_b6,
    efficientnet_b7,
)
from torchvision.models.googlenet import googlenet
from torchvision.models.inception import inception_v3
from torchvision.models.mnasnet import mnasnet0_5, mnasnet0_75, mnasnet1_0, mnasnet1_3
from torchvision.models.mobilenetv2 import mobilenet_v2
from torchvision.models.mobilenetv3 import mobilenet_v3_large, mobilenet_v3_small
from torchvision.models.regnet import (
    regnet_y_400mf,
    regnet_y_800mf,
    regnet_y_1_6gf,
    regnet_y_3_2gf,
    regnet_y_8gf,
    regnet_y_16gf,
    regnet_y_32gf,
    regnet_x_400mf,
    regnet_x_800mf,
    regnet_x_1_6gf,
    regnet_x_3_2gf,
    regnet_x_8gf,
    regnet_x_16gf,
    regnet_x_32gf,
)
from torchvision.models.resnet import (
    resnet18,
    resnet34,
    resnet50,
    resnet101,
    resnet152,
    resnext50_32x4d,
    resnext101_32x8d,
    wide_resnet50_2,
    wide_resnet101_2,
)

# segmentation
from torchvision.models.segmentation import (
    fcn_resnet50,
    fcn_resnet101,
    deeplabv3_resnet50,
    deeplabv3_resnet101,
    deeplabv3_mobilenet_v3_large,
    lraspp_mobilenet_v3_large,
)
from torchvision.models.shufflenetv2 import shufflenet_v2_x0_5, shufflenet_v2_x1_0
from torchvision.models.squeezenet import squeezenet1_0, squeezenet1_1
from torchvision.models.vgg import vgg11, vgg13, vgg16, vgg19, vgg11_bn, vgg13_bn, vgg16_bn, vgg19_bn