import torch import torch.nn as nn import torchvision import torchvision.transforms as transforms # 配置GPU或CPU设置 device = torch.device(‘cuda:0‘ if torch.cuda.is_available() else ‘cpu‘) # 超参数设置 num_epochs = 5 num_classes = 10 batch_size = 100 learning_rate = 0.001 # 下载 MNIST dataset train_dataset = torchvision.datasets.MNIST(root=‘./data/‘, train=True, transform=transforms.ToTensor(),# 将PIL Image或者 ndarray 转换为tensor,并且归一化至[0-1],归一化至[0-1]是直接除以255 download=True) test_dataset = torchvision.datasets.MNIST(root=‘./data/‘, train=False, transform=transforms.ToTensor())# 将PIL Image或者 ndarray 转换为tensor,并且归一化至[0-1],归一化至[0-1]是直接除以255 # 训练数据加载,按照batch_size大小加载,并随机打乱 train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True) # 测试数据加载,按照batch_size大小加载 test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False) # Convolutional neural network (two convolutional layers) 2层卷积 class ConvNet(nn.Module): def __init__(self, num_classes=10): super(ConvNet, self).__init__() self.layer1 = nn.Sequential( nn.Conv2d(1, 16, kernel_size=5, stride=1, padding=2), nn.BatchNorm2d(16), nn.ReLU(), nn.MaxPool2d(kernel_size=2, stride=2)) self.layer2 = nn.Sequential( nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=2), nn.BatchNorm2d(32), nn.ReLU(), nn.MaxPool2d(kernel_size=2, stride=2)) self.fc = nn.Linear(7 * 7 * 32, num_classes) def forward(self, x): out = self.layer1(x) out = self.layer2(out) out = out.reshape(out.size(0), -1) out = self.fc(out) return out model = ConvNet(num_classes).to(device) print(model) # ConvNet( # (layer1): Sequential( # (0): Conv2d(1, 16, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2)) # (1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) # (2): ReLU() # (3): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)) # (layer2): Sequential( # (0): Conv2d(16, 32, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2)) # (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) # (2): ReLU() # (3): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)) # (fc): Linear(in_features=1568, out_features=10, bias=True)) # 损失函数与优化器设置 # 损失函数 criterion = nn.CrossEntropyLoss() # 优化器设置 ,并传入CNN模型参数和相应的学习率 optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) # 训练CNN模型 total_step = len(train_loader) for epoch in range(num_epochs): for i, (images, labels) in enumerate(train_loader): images = images.to(device) labels = labels.to(device) # 前向传播 outputs = model(images) # 计算损失 loss loss = criterion(outputs, labels) # 反向传播与优化 # 清空上一步的残余更新参数值 optimizer.zero_grad() # 反向传播 loss.backward() # 将参数更新值施加到RNN model的parameters上 optimizer.step() # 每迭代一定步骤,打印结果值 if (i + 1) % 100 == 0: print (‘Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}‘ .format(epoch + 1, num_epochs, i + 1, total_step, loss.item())) # 测试模型 # model.train model.eval 在测试模型时在前面使用:model.eval() ; 在训练模型时会在前面加上:model.train() # 让model变成测试模式,是针对model 在训练时和评价时不同的 Batch Normalization 和 Dropout 方法模式 # eval()时,让model变成测试模式, pytorch会自动把BN和DropOut固定住,不会取平均,而是用训练好的值, # 不然的话,一旦test的batch_size过小,很容易就会被BN层导致生成图片颜色失真极大。 model.eval() # eval mode (batchnorm uses moving mean/variance instead of mini-batch mean/variance) with torch.no_grad(): correct = 0 total = 0 for images, labels in test_loader: images = images.to(device) labels = labels.to(device) outputs = model(images) _, predicted = torch.max(outputs.data, 1) total += labels.size(0) correct += (predicted == labels).sum().item() print(‘Test Accuracy of the model on the 10000 test images: {} %‘.format(100 * correct / total)) # 保存已经训练好的模型 # Save the model checkpoint torch.save(model.state_dict(), ‘model.ckpt‘)
原文地址:https://www.cnblogs.com/jeshy/p/11438484.html
时间: 2024-10-10 03:11:08