pytorch 完整的模型训练过程（入门级代码）

pytorch 模型训练过程（入门级代码）

加载数据集

Dataset是一个包装类，用来将数据包装为Dataset类，然后传入DataLoader中，我们再使用DataLoader来更加快捷的对数据进行操作。

Dataset

数据集包括训练集和测试集，在pytorch torchvision中有些常见的数据集，可以直接写代码设置参数下载。例如下载CIFAR10数据集：

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

download为True时会自动下载，下载的数据为PIL格式，要转为tensor。

每个数据集下载时的参数设置可能不一样，可以去官方文档具体查看。

如果是自己的数据集，大多数时候需要重写dataset类以方便对数据操作，入门级先不做介绍。

DataLoader

train_dataloader=DataLoader(dataset=train_data,batch_size=64,shuffle=True,num_works=0,drop_last=False)
test_dataloader=DataLoader(dataset=test_data,batch_size=64,shuffle=True,num_works=0,drop_last=False)

batch_size: 即每个batch的数量，默认1

shuffle： 是否打乱数据，默认False，常用True

num_works：加载数据集使用几个子进程

drop_last：最后一个batch时不满足设置数量时，是否舍弃

只介绍几个常用参数，只有第一个数据集参数是必选的，其他都有默认选项。DataLoader取出的数据包括数据和标签。

创建模型

模型一般会单独放在一个python文件中，使用时 import*

下面是一个CIFAR10 模型图

class CIFAR(nn.Module):
    def __init__(self):
        super(CIFAR, self).__init__()
        self.conv1 = Conv2d(3,32,5,padding=2)
        self.maxpool1 = MaxPool2d(2)
        self.conv2 = Conv2d(32,32,5,padding=2)
        self.maxpool2 = MaxPool2d(2)
        self.conv3 = Conv2d(32,64,5,padding=2)
        self.maxpool3 = MaxPool2d(2)
        self.flatten = Flatten()
        self.linear1 = Linear(64*4*4,64)
        self.linear2 = Linear(64,10)
    def forward(self,x):
        x = self.conv1(x)
        x = self.maxpool1(x)
        x = self.conv2(x)
        x = self.maxpool2(x)
        x = self.conv3(x)
        x = self.maxpool3(x)
        x = self.flatten(x)
        x = self.linear1(x)
        x = self.linear2(x)
        return x

按照官方文档的这个公式来计算padding和stride等。

训练和测试

全都只写了最简单的步骤以便理解，其实还有很多详细代码来可视化训练过程，以后慢慢学

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

#优化器
learning_rate = 0.01
optimizer = torch.optim.SGD(cifar.parameters(),lr = learning_rate)
#训练轮数
epoch = 10
for i in range(epoch):
	#训练
    cifar.train()
	for data in train_dataloader
        imgs,targets = data
        outputs = cifar(imgs)
        loss = loss_fn
        #梯度清零
        optimizer.zero_grad()
        #误差反向传播
        loss.backward()
        #参数更新
        optimizer.step()
    #测试
    cifar.eval()
    with torch.no_grad():
        for data in test_dataloader:
            imgs,targets = data
            outputs = cifar (imgs)
            loss = loss_fn(outputs,targets)      

GPU训练

要将模型和数据迁移到gpu上，常用代码

device = torch.device("cude" if torch.cuda.is_available() else "cpu")
model.to(device)

强烈建议看b站入门级教程，讲解非常仔细，本文参考该课程 课程地址