创建模型

模型创建简介

模型创建是构建图像分类模型的第一步。在本课中，我们将学习如何根据所需的架构创建一个新模型。

导入

在文件顶部添加以下导入语句，以引入 PyTorch 图像模型库（timm）：

1import timm

实例化新模型

模型实例化通过内置的 create_model(modelName) 函数完成。模型名称是必需参数，必须传入。现在我们可以像使用普通模型一样使用它。为了理解 PyTorch 模型的实际工作原理，请研究以下组件中的代码：

1import timm 
2
3model = timm.create_model('resnet50', pretrained=False, checkpoint_path="/app/resnet50_best.pth.tar", num_classes=4)
4model.eval()
5
6# 获取模型中特征的总数
7print('该数据集中特征的数量:', model.num_features)
8
9# 获取模型中类别的总数
10print('该数据集中类别的数量:', model.num_classes)

参数说明

create_model 函数接受以下参数：

• model_name (str)：要实例化的模型名称。
• pretrained (bool)：若为 True，则加载预训练的 ImageNet-1k 权重。
• checkpoint_path (str)：模型初始化后加载检查点的路径。

如果我们使用的是自定义训练的模型，只需将 checkpoint_path 的值更改为对应模型的路径即可。

其他关键字参数

此外，还可以传入其他特定于模型的参数。所有模型共享以下通用参数：

• drop_rate (float)：训练时的 Dropout 率（默认值：0.0）。
• global_pool (str)：全局池化类型（默认值：'avg'）。

对于 resnet50 模型还需要以下额外参数：

• num_classes (int)：可用类别数量。
• in_chans (int)：图像数据集的通道数。

请参考以下示例：

1import timm
2
3model = timm.create_model(
4    "resnet50",
5    num_classes=4,
6    in_chans=3,
7    pretrained=False,
8    checkpoint_path="/app/resnet50_best.pth.tar"
9)
10model.eval()
11
12print('resnet50 的 Dropout 率:', model.drop_rate)
13print('resnet50 的全局池化方式:', model.global_pool)

注意：SelectAdaptivePool2d 类对输入应用二维自适应池化操作。

PyTorch 图像模型提供了一个统一的模型创建函数。我们可以通过 create_model(modelName) 来实例化一个新模型。