pytorch_demo1

一个简单的分类模型
model.python

import torchvision
import torch

# 准备数据集
from torch import nn
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,1,2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 32, 5, 1, 2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 5, 1, 2),
            nn.MaxPool2d(2),
            nn.Flatten(),
            nn.Linear(64*4*4,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)

train.python

import torchvision
import torch

# 准备数据集
from torch import nn
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,1,2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 32, 5, 1, 2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 5, 1, 2),
            nn.MaxPool2d(2),
            nn.Flatten(),
            nn.Linear(64*4*4,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)

运行：在train.py文件中运行
在terminal中输入：tensorboard --logdir=logs/train_logs(自己的日志存在哪，这里就填哪)
效果如图：