UP:B站-刘二大人
原视频链接:10.卷积神经网络(基础篇)_哔哩哔哩_bilibili
import torch
from torchvision import datasets
from torchvision import transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
batch_size = 64
transform = transforms.Compose( [transforms.ToTensor(), transforms.Normalize ((0.1307, ), (0.3081, )) ] )
train_dataset = datasets.MNIST(root='../dataset/mnist/', train=True, download=True, transform=transform)
train_loader = DataLoader(dataset=train_dataset, shuffle=True, batch_size=batch_size)
test_dataset = datasets.MNIST(root='../dataset/mnist/', train=False, download=True, transform=transform)
test_loader = DataLoader(dataset=test_dataset, shuffle=False, batch_size=batch_size)
class Model(torch.nn.Module):
def __init__(self):
super(Model, self).__init__()
self.conv1 = torch.nn.Conv2d(1, 10, kernel_size = 5) # 三个参数分别是:输入的通道数,输出的通道数,卷积核的大小(5×5)
self.conv2 = torch.nn.Conv2d(10, 20, kernel_size = 3, padding=1)
self.conv3 = torch.nn.Conv2d(20, 32, kernel_size = 3)
self.pooling = torch.nn.MaxPool2d(2) # 2*2的池化窗口
self.active = torch.nn.ReLU()
self.fc1 = torch.nn.Linear(128, 64) # 全连接层
self.fc2 = torch.nn.Linear(64, 32)
self.fc3 = torch.nn.Linear(32, 10)
def forward(self, x):
x = self.pooling(self.active(self.conv1(x)))
x = self.pooling(self.active(self.conv2(x)))
x = self.pooling(self.active(self.conv3(x)))
# 此时的x是维度为4的tensor,即(batch_size, c , H, W),x.size(0)指的是batch_size的值
x = x.view(x.size(0), -1)
x = self.fc1(x)
x = self.fc2(x)
x = self.fc3(x)
return x
model = Model()
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.5)
loss_list = []
def train(epoch):
loss_sum = 0.0
for i, (inputs, labels) in enumerate(train_loader):
outputs = model(inputs)
loss = criterion(outputs, labels)
loss_sum += loss.item()
loss_list.append(loss.item())
optimizer.zero_grad()
loss.backward()
optimizer.step()
if i % 300 ==299:
print("[%d %5d] loss:%.3f" % (epoch+1, i+1, loss_sum / i ))
#if __name__=="__main__":
#for epoch in range(10):
#train(epoch)
def test():
correct = 0
total = 0
with torch.no_grad(): # 测试不用计算梯度
for (images, targets) in test_loader:
y_pred = model(images)
_, predicted = torch.max(y_pred.data, dim=1)
total += targets.size(0)
correct += (predicted == targets).sum().item()
print("Accuarcy on test dataset:%.3f %% [%d %d]"%(100*correct/total, correct, total))
if __name__=="__main__":
for epoch in range(5):
train(epoch)
test()
plt.plot(range(len(loss_list)), loss_list)
plt.xlabel("step")
plt.ylabel("loss")
plt.show()