model.py

原创已于 2022-04-12 10:56:29 修改 · 101 阅读

0 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch

于 2022-04-12 10:53:45 首次发布

深度学习专栏收录该内容

51 篇文章

订阅专栏

该博客展示了如何使用PyTorch构建一个基本的卷积神经网络（CNN）模型，包括多个卷积层、最大池化层和全连接层。模型用于图像分类任务，输入为3通道32x32像素图像，最后输出为10类别的概率分布。

部署运行你感兴趣的模型镜像

import torch
import torchvision
from torch import nn
from torch.nn import Conv2d, MaxPool2d, Linear
from torch.nn.modules.flatten import Flatten
from torch.utils.data import DataLoader
class Tudui(nn.Module):
    def __init__(self):
        super(Tudui, self).__init__()
        self.model = nn.Sequential(
            nn.Conv2d(3, 32, 5, padding=2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 32, 5, padding=2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 5, padding=2),
            nn.MaxPool2d(2),
            nn.Flatten(),
            nn.Linear(1024, 64),
            nn.Linear(64, 10)
        )
    def forward(self, x):
        x = self.model(x)
        return x

if __name__ == '__main__':
    tudui = Tudui()
    input = torch.ones((64, 3, 32, 32))
    output = tudui(input)
    print(output.shape)

您可能感兴趣的与本文相关的镜像