pytorch模型训练全过程

#model.py
#搭建神经网络
class Model_test(nn.Module):
	def __init__():
		super(Model_test, self).__init__()
		self.model = nn.Sequential(
		nn.Conv2d(3,32,5,1,2),
		nn.MaxPool2d(2),
		nn.Conv2d(32,32,5,1,2),
		nn.MaxPool2d(2),
		nn.Conv2d(32,64,5,1,2),
		nn.MaxPool2d(2),
		nn.Flatten(),
		nn,Linear(64*4*4,64),
		nn,Linear(64,10))
	def forward(self,x):
		X= self.model(x)
		return x
#测试网络正确性
if __name__=="__main__"
	model_test = Model_test()
	input = torch.ones((64,3,32,32))
	output = model_test(input)
	print(output.shape)
	
	

#train.py

import torchvision
from model import * #model.py和train.py需要在同一个文件夹下
import time  #计时


#定义训练的设备
device = torch.device("cpu")
device = torch.device("cuda")
device = torch.device("cuda:0")
#准备数据集
train_data = torchvision.datasets.CIFAR10("./data", train= Ture, transform=torchvision.transforms.ToTensor(),download=True)
test_data = torchvision.datasets.CIFAR10("./data", train= False, transform=torchvision.transforms.ToTensor(),download=True)

#length长度
train_data_size = len(train_data)
test_data_size = len(test_data)
#如果train_data_size=10，训练数据集的长度为：10
print("训练数据集的长度为：{}".format(train_data_size))
print("测试数据集的长度为：{}".format(test_data_size))

#利用DataLoader来加载数据集
train_dataloader = DataLoader(train_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)

#创建网络模型
model_test = Model_test()
#if torch.cuda,is_available():
#	model_test.cuda()  ####方法一GPU 1.网络模型
model_test.to(device)  ###方法二GPU

#损失函数
loss_fn= nn.CrossEntropyLoss()
if torch.cuda,is_available():
	loss_fn = loss_fn.cuda()  ####方法一GPU  2.损失函数
loss_fn = loss_fn.to(device)  ###方法二GPU

#优化器
#1e-2 = 1x(10)^(-2)=0.01
learning_rate = 1e-2
optimizer = torch.optim.SGD(model_test.paaramenters(), lr = learning_rate)


#设置训练网络的一些参数
total_train_step = 0
#记录测试的次数
total_test_step = 0
#训练的轮数
epoch = 10

#添加tensorboard
writer = SummaryWriter("logs_train")
start_time = time.time()


for i in range(epoch):
	print("-------第{}轮的训练开始了-------".format(i+1))
	#训练开始
	model_test.train()   #作用对特定层起作用
	for data in train_dataloader:
		imgs, targets = data
		#if torch.cuda,is_available():
		#	imgs = imgs.cuda()   ####方法一GPU  3.1 训练数据
		#	targets = targets.cuda()   ####方法一GPU  3.1 训练数据
		imgs = imgs.to(device)  ###方法二GPU
		targets = targets.to(device)  ###方法二GPU
		outputs = model_test(imgs)
		loss = loss_fn(outputs, targets)
		#优化器优化模型
		optimizer.zero_grad()
		loss.backward()
		optimizer.step()
		total_train_step +=1
		if total_train_step %100 =0:
			end_time = time.time()
			print(start_tim, end_time)
			print("训练次数：{}，loss：{}".format(total_train_step,loss.item()))
			writer.add_scalar("train_lose",loss.item(), total_train_step)
	
	#测试步骤开始
	model_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()   ####方法一GPU  3.2 测试数据
			#	targets = targets.cuda()   ####方法一GPU  3.2 测试数据
			imgs = imgs.to(device)  ###方法二GPU
			targets = targets.to(device)  ###方法二GPU
			outputs = model_test(imgs)
			loss = loss_fn(outputs, targets)
			total_test_loss = total_test_loss + loss
			accuracy = (output.agxmax(1)==taargets).sum
			total_accuracy = total_accuracy + accuracy
	print("整体测试集的Loss:{}".format(total_test_loss))
	print("整体测试集的准确率:{}".format(total_accuracy、test_data_size))
	writer.add_scalar("test_lose",total_test_loss.item(), total_test_step)
	total_test_step +=1
	torch.save(model_test, "model_test_{}.pth".format(i))
	#torch.save(“Model_test_{}.pth”.format(i))
	print("模型已保存")
writer.close()
		
	

test2.py
outputs = torcch.tensor([0.3,0.4],[0.1,0.2])
print(outputs.argmax(1)) #0代表纵向，1代表横向

preds = outputs.argmax(1)
targets = torch.tensor([0,1])
print((preds == targets).sum())

模型在gpu上训练的三个关键：网络模型、数据（输入、标注）、损失函数
调用.cuda()

#train_gpu.py

方法一：.cuda()
方法二：.to(device)
device = torch.device
torch.device(“cuda:0”)
torch.device(“cuda:1”)
方法三：
device = torch.device(“cuda” if torch.cuda.is_available() else “cpu”)