import torch
from torchvision import transforms # 是一个常用的图片变换类
from torchvision import datasets
from torch.utils.data import DataLoader
import torch.nn.functional as F
batch_size = 64
transform = transforms.Compose(
[
transforms.ToTensor(), # 把数据转换成张量
transforms.Normalize((0.1307,), (0.3081,)) # 0.1307是均值，0.3081是标准差
]
)
train_dataset = datasets.MNIST(root='.\mnist',
train=True,
download=False,
transform=transform)
train_loader = DataLoader(train_dataset,
shuffle=True,
batch_size=batch_size)
test_dataset = datasets.MNIST(root='.\mnist',
train=False,
download=False,
transform=transform)
test_loader = DataLoader(test_dataset,
shuffle=True,
batch_size=batch_size)
class CNN(torch.nn.Module):
def __init__(self):
super(CNN, self).__init__()
self.layer1 = torch.nn.Sequential(
torch.nn.Conv2d(1, 25, kernel_size=3),
torch.nn.BatchNorm2d(25),
torch.nn.ReLU(inplace=True)
)
self.layer2 = torch.nn.Sequential(
torch.nn.MaxPool2d(kernel_size=2, stride=2)
)
self.layer3 = torch.nn.Sequential(
torch.nn.Conv2d(25, 50, kernel_size=3),
torch.nn.BatchNorm2d(50),
torch.nn.ReLU(inplace=True)
)
self.layer4 = torch.nn.Sequential(
torch.nn.MaxPool2d(kernel_size=2, stride=2)
)
self.fc = torch.nn.Sequential(
torch.nn.Linear(50 * 5 * 5, 1024),
torch.nn.ReLU(inplace=True),
torch.nn.Linear(1024, 128),
torch.nn.ReLU(inplace=True),
torch.nn.Linear(128, 10)
)
def forward(self, x):
# print("forward")
x = self.layer1(x)
# print(x.shape)
x = self.layer2(x)
# print(x.shape)
x = self.layer3(x)
# print(x.shape)
x = self.layer4(x)
# print(x.shape)
x = x.view(x.size(0), -1) # 在进入全连接层之前需要把数据拉直Flatten
x = self.fc(x)
return x
model = CNN()
# 下面两行代码主要是如果有GPU那么就使用GPU跑代码，否则就使用cpu。cuda:0表示第1块显卡
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") # 将数据放在GPU上跑所需要的代码
model.to(device) # 将数据放在GPU上跑所需要的代码
criterion = torch.nn.CrossEntropyLoss() # 使用交叉熵损失
optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.5) # momentum表示冲量，冲出局部最小
def train(epochs):
running_loss = 0.0
for batch_idx, data in enumerate(train_loader, 0):
inputs, target = data
# print(inputs.shape,target.shape)
inputs, target = inputs.to(device), target.to(device) # 将数据放在GPU上跑所需要的代码
# print(inputs.shape)
optimizer.zero_grad()
# 前馈+反馈+更新
outputs = model(inputs)
loss = criterion(outputs, target)
loss.backward()
optimizer.step()
running_loss += loss.item()
# if batch_idx % 10 == 0: # 不让他每一次小的迭代就输出，而是300次小迭代再输出一次
# print('[%d,%5d] loss:%.3f' % (epoch + 1, batch_idx + 1, running_loss / 10))
# running_loss = 0.0
torch.save(model, 'model_{}.pth'.format(epochs))
def test():
correct = 0
total = 0
with torch.no_grad(): # 下面的代码就不会再计算梯度
for data in test_loader:
inputs, target = data
inputs, target = inputs.to(device), target.to(device) # 将数据放在GPU上跑所需要的代码
outputs = model(inputs)
_, predicted = torch.max(outputs.data, dim=1) # _为每一行的最大值，predicted表示每一行最大值的下标
total += target.size(0)
correct += (predicted == target).sum().item()
print('Accuracy on test set:%d %%' % (100 * correct / total))
if __name__ == '__main__':
for epoch in range(10):
train(epoch)
test()
