Pytorch学习笔记(五)
一、如何使用GPU训练
方法一:在网络模型,数据和损失函数处调用cuda()方法
在上一节的模型训练套路的代码中直接进行修改
#准备数据集
train_data = torchvision.datasets.CIFAR10(root="./dataset", train=True, transform=torchvision.transforms.ToTensor(), download=True)
test_data = torchvision.datasets.CIFAR10(root="./dataset", train=False, transform=torchvision.transforms.ToTensor(), download=True)
#加载数据集
train_dataloader = DataLoader(train_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)
#搭建神经网络,可以将这一部分的代码单独放到一个model的python中,model文件应该和本文件在一个文件夹下
class Test(nn.Module):
def __init__(self):
super(Test, self).__init__()
#使用Sequential提高代码的简洁可读性
self.model1 = Sequential{
Conv2d(3, 32, 5, padding=2),
MaxPool2d(2),
Conv2d(32, 64, 5, padding=2),
MaxPool2d(2),
Conv2d(32, 64, 5, padding=2),
MaxPool2d(2),
Flatten(),
Linear(1024, 64),
Linear(64, 10)
)
def forward(self, x):
x = self.model1(x)
return x
#创建网络模型
test = Test()
if torch.cuda.is_available():
test = test.cuda()
#损失函数
loss_fn = nn.CrossEntropyLoss()
if torch.cuda.is_available():
loss_fn = loss_fn.cuda()
#优化器
learning_rate = 0.01 #或为1e-2
optimizer = torch.optim.SGD(test.parameters(), lr = learning_rate)
#设置训练网络的参数
total_train_step = 0 #记录训练的次数
total_test_step = 0 #记录测试的次数
epoch = 10 #训练的迭代次数
for i in range(epoch):
print("---------第{}轮训练开始----------".format(i+1))
#训练步骤开始
test.train() #只对某些特定的层有作用
for data in train_dataloader:
imgs, targets = data
if torch.cuda.is_available():
imgs = imgs.cuda()
targets = targets.cuda()
outputs = test(imgs)
loss = loss_fn(outputs, targets)
#优化器优化模型
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_train_setp = total_train_step + 1
if total_train_step % 100 == 0
print("训练次数:{}, loss:{}".format(total_train_step, loss.item()))
#测试步骤开始
test.eval() #只对某些特定的层有作用
total_test_loss = 0
total_accuracy = 0
with torch.no_grad():
for data in test_dataloader:
imgs, targets = data
if torch.cuda.is_available():
imgs = imgs.cuda()
targets = targets.cuda()
outputs = test(imgs)
loss = loss_fn(outputs, targets)
total_test_loss = total_test_loss + loss.item()
accuracy = (outputs.argmax(1) == targets).sum() #argmax将输出的概率值转化成1/0
total_accuacy = total_accuracy + accuracy
print("整体测试集上的Loss:{}".format(total_test_loss))
print("整体测试集上的正确率:{}".format(total_accuracy/len(test_data)))
torch.save(test, "test_{}.pth".format(i))
方法二:定义运行设备,使用时直接修改device
#定义训练的设备
device = torch.device("cuda")
#准备数据集
train_data = torchvision.datasets.CIFAR10(root="./dataset", train=True, transform=torchvision.transforms.ToTensor(), download=True)
test_data = torchvision.datasets.CIFAR10(root="./dataset", train=False, transform=torchvision.transforms.ToTensor(), download=True)
#加载数据集
train_dataloader = DataLoader(train_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)
#搭建神经网络,可以将这一部分的代码单独放到一个model的python中,model文件应该和本文件在一个文件夹下
class Test(nn.Module):
def __init__(self):
super(Test, self).__init__()
#使用Sequential提高代码的简洁可读性
self.model1 = Sequential{
Conv2d(3, 32, 5, padding=2),
MaxPool2d(2),
Conv2d(32, 64, 5, padding=2),
MaxPool2d(2),
Conv2d(32, 64, 5, padding=2),
MaxPool2d(2),
Flatten(),
Linear(1024, 64),
Linear(64, 10)
)
def forward(self, x):
x = self.model1(x)
return x
#创建网络模型
test = Test()
test.to(device)
#损失函数
loss_fn = nn.CrossEntropyLoss()
loss_fn = loss_fn.to(device)
#优化器
learning_rate = 0.01 #或为1e-2
optimizer = torch.optim.SGD(test.parameters(), lr = learning_rate)
#设置训练网络的参数
total_train_step = 0 #记录训练的次数
total_test_step = 0 #记录测试的次数
epoch = 10 #训练的迭代次数
for i in range(epoch):
print("---------第{}轮训练开始----------".format(i+1))
#训练步骤开始
test.train() #只对某些特定的层有作用
for data in train_dataloader:
imgs, targets = data
imgs = imgs.to(device)
targets = targets.to(device)
outputs = test(imgs)
loss = loss_fn(outputs, targets)
#优化器优化模型
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_train_setp = total_train_step + 1
if total_train_step % 100 == 0
print("训练次数:{}, loss:{}".format(total_train_step, loss.item()))
#测试步骤开始
test.eval() #只对某些特定的层有作用
total_test_loss = 0
total_accuracy = 0
with torch.no_grad():
for data in test_dataloader:
imgs, targets = data
imgs = imgs.to(device)
targets = targets.to(de)
outputs = test(imgs)
loss = loss_fn(outputs, targets)
total_test_loss = total_test_loss + loss.item()
accuracy = (outputs.argmax(1) == targets).sum() #argmax将输出的概率值转化成1/0
total_accuacy = total_accuracy + accuracy
print("整体测试集上的Loss:{}".format(total_test_loss))
print("整体测试集上的正确率:{}".format(total_accuracy/len(test_data)))
torch.save(test, "test_{}.pth".format(i))
完整的模型测试套路
利用已经训练好的模型,然后给他提供输入
model.eval()
with torch.no_grad():
output = model(image) #注意这里的图片是四维的tensor数据类型
print(output)
print(output.argmax(1))
结束
完结!大家可以去GitHub上面下载一些pytorch的代码训练学习啦!