LeNet和CIFAR10数据集为例解析Pytorch常用操作

一.搭建神经网络

import torch.nn as nn
import torch.nn.functional as F

# 继承了nn.Module类
class LeNet(nn.Module):
def __init__(self):
super(LeNet, self).__init__()
self.conv1 = nn.Conv2d(3, 6, 5)
self.conv2 = nn.Conv2d(6, 16, 5)
self.fc1 = nn.Linear(400, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10)

def forward(self, x):
out = F.relu(self.conv1(x))
out = F.max_pool2d(out, 2)
out = F.relu(self.conv2(out))
out = F.max_pool2d(out, 2)
out = out.view(out.size(0), -1)
out = F.relu(self.fc1(out))
out = F.relu(self.fc2(out))
out = self.fc3(out)
return out


二,应用搭建的DNN

import argparse
parser = argparse.ArgumentParser(description='PyTorch CIFAR10 Training')
parser.add_argument('--resume', '-r', action='store_true', help='resume from checkpoint')
args = parser.parse_args()
# 使用args.lr,args.resume可以获得参数


import torchvision
import torchvision.transforms as transforms

cifar_norm_mean = (0.49139968, 0.48215827, 0.44653124)
cifar_norm_std = (0.24703233, 0.24348505, 0.26158768)
transform_train = transforms.Compose([
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(cifar_norm_mean, cifar_norm_std),
])

transform_test = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(cifar_norm_mean, cifar_norm_std),
])

trainset = torchvision.datasets.CIFAR10\
(trainset, batch_size=128, shuffle=True, num_workers=2)

testset = torchvision.datasets.CIFAR10\
(testset, batch_size=100, shuffle=False, num_workers=2)

classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')



import torch
import lenet
import torch.backends.cudnn as cudnn

device = 'cuda' if torch.cuda.is_available() else 'cpu'
net = lenet.LeNet()
net = net.to(device) # 将net基于设备运行
if device == 'cuda':
net = torch.nn.DataParallel(net) # 单机多GPUs的并行处理,只有1个GPU这句话没用
cudnn.benchmark = True # 将选择最佳优化算法,加快训练速度


import torch
import torch.nn as nn
import torch.optim as optim

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=0.9, weight_decay=5e-4)

# Training
def train(epoch):
print('\nEpoch: %d' % epoch)
net.train()
train_loss = 0
correct = 0
total = 0
# 遍历训练集, 梯度下降训练的框架
for batch_idx, (inputs, targets) in enumerate(trainloader):
inputs, targets = inputs.to(device), targets.to(device)
outputs = net(inputs)  # 将训练数据输入网络得到输出
loss = criterion(outputs, targets) # 用定义的交叉熵函数对象计算误差
loss.backward() # 计算反向梯度
optimizer.step() # 利用计算好的梯度更新参数

# 输出loss和预测信息
train_loss += loss.item() # 将tensor([0.98])变为标量0.98
_, predicted = outputs.max(1)  # 1表示对第一维约简,_为返回的最大值,predicted为最大值的下标
total += targets.size(0) # .size返回一个元组表示维度
correct += predicted.eq(targets).sum().item() # .eq返回一个原维度的bool tensor
print('Train accuracy:', correct/total)


def test(epoch):
global best_acc
net.eval() # 将net设为评估模式,对dropout和batchnorm将变为评估模式
test_loss = 0
correct = 0
total = 0
for batch_idx, (inputs, targets) in enumerate(testloader):
inputs, targets = inputs.to(device), targets.to(device)
outputs = net(inputs)
loss = criterion(outputs, targets)

test_loss += loss.item()
_, predicted = outputs.max(1)
total += targets.size(0)
correct += predicted.eq(targets).sum().item()
print('Test accuracy:', correct/total)


参数保存与网络断点的设置

state = {
'net': net.state_dict(),  # 获取当前net的参数字典,为OrderedDict
'acc': acc,  # 当前准确率
'epoch': epoch, # 当前轮数
'optimizer': optimizer.state_dict()
}
if not os.path.isdir('checkpoint'):
os.mkdir('checkpoint')
# 保存状态字典到pkl文件
torch.save(state, './checkpoint/ckpt.pkl')
# 也可以直接将参数字典net.state_dict()保存到pkl文件
torch.save(state, './checkpoint/params.pkl')


checkpoint = torch.load('./checkpoint/ckpt.pkl')
best_acc = checkpoint['acc']
start_epoch = checkpoint['epoch']


03-11 1万+
03-21 1149
04-04 2万+
04-15 297
04-25 1818
03-12 630
11-22 404