基于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()

执行代码后输出:

torch.Size([5, 2])

你可能感兴趣的:(Pytorch,图像分类,深度学习,神经网络,深度学习,python,ResNet,PyTorch)