完整的模型训练套路

train.py
import torchvision
from model import *
# 准备数据集
from torch.utils.data import DataLoader

train_data = torchvision.datasets.CIFAR10(root="../data", train=True,
transform = torchvision.transforms.ToTensor(),download=True)
test_data = torchvision.datasets.CIFAR10(root="../data",train=False,
transform=torchvision.transforms.TopTensor(),download=True)

# length 长度
train_data_size = len(train_data)
test_data_size = len(test_data)
p rint("训练数据集的长度为:{}".format(train_data_size))
print("测试数据集的长度为:{}".format(test_data_size))

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

# 创建网络模型
cifar = CIFar()

# 损失函数
loss_fn= nn.CrossEntropyLoss()

# 优化器
learning_ratr = 0.01
optimizer = torch.optim.SGD(cifar.parameters(),lr=learning_ratr)

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

for i in range(epoch)
	print("第{}轮训练开始".format(i+1))

	#训练步骤开始
	cifar.train()
	for data in train_dataloader:
		imgs, target = data
		outputs = cifar(imgs)
		loss = loss_fn(outputs,targets)

		#优化器优化模型
		optimizer.zero_grad()
		loss.backward()
		optimizer.step()

		total_train_step = total_train_step + 1
		print("训练次数:{},loss:{}".format(total_train_step,loss.item()))

	# 测试步骤开始
	cifar.eval()
	total_test_loss = 0
	with torch.no_grad():
		for data in test_dataloader:
			imgs,targets = data
			ouyputs = cifar(imgs)
			loss = loss_fn(outputs,targets)
			total_test_loss = total_test_loss + loss.item()
	print("整体测试集上的Loss：{}"。foramt(total_test_loss))

	
		

model.py
# 搭建神经网络
class CIFar(nn.Module):
	 def __init__(self):
	 	super(CIFar,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__':
    cifar = CIFar()
    input = torch.ones((64,3,32,32))
    output = CIFar(input)
    print(output.shap)