一个简单的分类模型
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(自己的日志存在哪,这里就填哪)
效果如图: