1.LeNet网络
import torch
from torch import nn, optim
import torch.utils.data as Data
import torchvision.transforms as transforms
import torchvision
import sys
import time
mnist_train = torchvision.datasets.FashionMNIST(root=".", train=True, download=True,
transform=transforms.ToTensor())
mnist_test = torchvision.datasets.FashionMNIST(root=".", train=False, download=True,
transform=transforms.ToTensor())
batch_size = 256
if sys.platform.startswith("win"):
num_workers = 0
else:
num_workers = 8
train_iter = Data.DataLoader(mnist_train, batch_size=batch_size, shuffle=True, num_workers=num_workers)
test_iter = Data.DataLoader(mnist_test, batch_size=batch_size, shuffle=False, num_workers=num_workers)
class LeNet(nn.Module):
def __init__(self):
super(LeNet, self).__init__()
self.c1 = nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5)
self.a1 = nn.Sigmoid()
self.p1 = nn.MaxPool2d(kernel_size=2, stride=2)
self.c2 = nn.Conv2d(6, 16, 5)
self.a2 = nn.Sigmoid()
self.p2 = nn.MaxPool2d(2, 2)
self.l1 = nn.Linear(16*4*4, 120)
self.a3 = nn.Sigmoid()
self.l2 = nn.Linear(120, 84)
self.a4 = nn.Sigmoid()
self.l3 = nn.Linear(84, 10)
def forward(self, x):
x = self.c1(x)
x = self.a1(x)
x = self.p1(x)
x = self.c2(x)
x = self.a2(x)
x = self.p2(x)
x = x.view(x.shape[0], -1)
x = self.l1(x)
x = self.a3(x)
x = self.l2(x)
x = self.a4(x)
y = self.l3(x)
return y
net = LeNet()
net = net.cuda(0)
loss = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(), lr=0.01)
start = time.time()
for epoch in range(10):
for X, y in train_iter:
X, y = X.cuda(), y.cuda()
out = net(X)
l = loss(out, y)
optimizer.zero_grad()
l.backward()
optimizer.step()
print("epoch: %d loss: %f" % (epoch, l.item()))
print("time: %f" % (time.time()-start))