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的代码训练学习啦！