CIFAR10模型如下图(直接从网上搜的):
具体实现代码:
代码使用了一个Sequential序列,可以使得代码更简洁。
import torch
from torch import nn
from torch.utils.tensorboard import SummaryWriter
class Tudui(nn.Module):
def __init__(self):
super(Tudui, self).__init__()
self.model1 = nn.Sequential(
nn.Conv2d(in_channels=3, out_channels=32, kernel_size=5, padding=2),
nn.MaxPool2d(kernel_size=2),
nn.Conv2d(in_channels=32, out_channels=32, kernel_size=5, padding=2),
nn.MaxPool2d(kernel_size=2),
nn.Conv2d(in_channels=32, out_channels=64, kernel_size=5, padding=2),
nn.MaxPool2d(kernel_size=2),
# 展现数据
nn.Flatten(),
nn.Linear(in_features=1024, out_features=64),
nn.Linear(in_features=64, out_features=10)
)
def forward(self, x):
x = self.model1(x)
return x
tudui = Tudui()
print(tudui)
input = torch.ones(64, 3, 32, 32)
#tensorboard 可视化
writer = SummaryWriter("./logs_seq")
writer.add_graph(tudui,input)
writer.close()
tensorboard可视化结果: