基于PyTorch的卷积神经网络图像分类——猫狗大战（二）：使用Pytorch定义网络模型

文章目录

• 1. 需要用到的库
• 2. 模型定义
• 3. 测试

基于上一篇文章 https://blog.csdn.net/linghu8812/article/details/100044971，这次介绍一下网络模型的定义。

1. 需要用到的库

import torch
import torch.nn as nn
from torchvision.models import resnet50
           

基于ResNet50定义网络模型，在此基础上进行微调。

2. 模型定义

class CatVSDogNet(nn.Module):
    def __init__(self, num_class=2):
        super(CatVSDogNet, self).__init__()
        self.model = resnet50(pretrained=True)
        self.model.fc = nn.Linear(2048, num_class)

    def forward(self, x):
        output = self.model(x)
        return output
           

基于ImageNet进行初始化，ImageNet有1000类，将网络模型的全连接层修改为两类。

3. 测试

测试代码如下：

def main():
    net = CatVSDogNet()
    y = net(torch.randn(5, 3, 448, 448))
    print(y.size())


if __name__ == '__main__':
    main()
           

执行代码后输出：