import torch import torch.nn as nn import torch.nn.functional as F import torch.optim as optim from torchvision import datasets, transforms import matplotlib.pyplot as plt

定义LeNet_ReLU网络

class LeNet_ReLU(nn.Module): def init(self): super(LeNet_ReLU, self).init() self.conv1 = nn.Conv2d(1, 6, 5) self.pool1 = nn.MaxPool2d(2, 2) self.conv2 = nn.Conv2d(6, 16, 5) self.pool2 = nn.MaxPool2d(2, 2) self.fc1 = nn.Linear(16 * 4 * 4, 120) self.fc2 = nn.Linear(120, 84) self.fc3 = nn.Linear(84, 10)

def forward(self, x):
    x = self.pool1(F.relu(self.conv1(x)))
    x = self.pool2(F.relu(self.conv2(x)))
    x = x.view(-1, 16 * 4 * 4)
    x = F.relu(self.fc1(x))
    x = F.relu(self.fc2(x))
    x = self.fc3(x)
    return x

加载FashionMNIST数据集

transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))]) train_set = datasets.FashionMNIST('./data', download=True, train=True, transform=transform) test_set = datasets.FashionMNIST('./data', download=True, train=False, transform=transform) train_loader = torch.utils.data.DataLoader(train_set, batch_size=64, shuffle=True) test_loader = torch.utils.data.DataLoader(test_set, batch_size=64, shuffle=True)

训练模型

def train(model, device, train_loader, optimizer, epoch): model.train() for batch_idx, (data, target) in enumerate(train_loader): data, target = data.to(device), target.to(device) optimizer.zero_grad() output = model(data) loss = F.cross_entropy(output, target) loss.backward() optimizer.step() if batch_idx % 100 == 0: print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( epoch, batch_idx * len(data), len(train_loader.dataset), 100. * batch_idx / len(train_loader), loss.item()))

测试模型

def test(model, device, test_loader): model.eval() test_loss = 0 correct = 0 with torch.no_grad(): for data, target in test_loader: data, target = data.to(device), target.to(device) output = model(data) test_loss += F.cross_entropy(output, target, reduction='sum').item() # 将一批的损失相加 pred = output.argmax(dim=1, keepdim=True) # 找到概率最大的下标 correct += pred.eq(target.view_as(pred)).sum().item()

test_loss /= len(test_loader.dataset)
print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
    test_loss, correct, len(test_loader.dataset),
    100. * correct / len(test_loader.dataset)))
return test_loss, 100. * correct / len(test_loader.dataset)

主函数

def main(): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = LeNet_ReLU().to(device) optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.5)

train_loss_list = []
test_loss_list = []
train_acc_list = []
test_acc_list = []

for epoch in range(1, 11):
    train(model, device, train_loader, optimizer, epoch)
    train_loss, train_acc = test(model, device, train_loader)
    test_loss, test_acc = test(model, device, test_loader)
    train_loss_list.append(train_loss)
    test_loss_list.append(test_loss)
    train_acc_list.append(train_acc)
    test_acc_list.append(test_acc)

# 绘制损失函数曲线
plt.figure()
plt.plot(range(1, 11), train_loss_list, label='Train Loss')
plt.plot(range(1, 11), test_loss_list, label='Test Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.show()

# 绘制分类正确率曲线
plt.figure()
plt.plot(range(1, 11), train_acc_list, label='Train Accuracy')
plt.plot(range(1, 11), test_acc_list, label='Test Accuracy')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.legend()
plt.show()

if name == 'main': main(