一个完整的 Pytorch 深度学习项目代码,项目结构是怎样的?

旺旺小小超

总结一下自己使用 pytorch 写深度学习模型的心得,所有的 pytorch 模型都离不开下面的几大组件。

Network
...

创建一个 Network 类,继承 torch.nn.Module,在构造函数中用初始化成员变量为具体的网络层,在 forward 函数中使用成员变量搭建网络架构,模型的使用过程中 pytorch 会自动调用 forword 进行参数的前向传播,构建计算图。以下拿一个简单的 CNN 图像分类模型举例

class Network(nn.Module):
    def __init__(self):
        super(Network, self).__init__()
        # 灰度图像的channels=1即in_channels=1 输出为10个类别即out_features=10
        # parameter(形参)=argument(实参) 卷积核即卷积滤波器 out_channels=6即6个卷积核 输出6个feature-maps(特征映射)
        # 权重shape 6*1*5*5
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5)
        self.bn1 = nn.BatchNorm2d(6)  # 二维批归一化 输入size=6
        # 权重shape 12*1*5*5
        self.conv2 = nn.Conv2d(in_channels=6, out_channels=12, kernel_size=5)
        
        # 全连接层:fc or dense or linear out_features即特征(一阶张量)
        # 权重shape 120*192
        self.fc1 = nn.Linear(in_features=12*4*4, out_features=120)
        self.bn2 = nn.BatchNorm1d(120)  # 一维批归一化 输入size=120
        # 权重shape 60*120
        self.fc2 = nn.Linear(in_features=120, out_features=60)
        # 权重shape 10*60
        self.out = nn.Linear(in_features=60, out_features=10)
        
    def forward(self, t):
        # (1) input layer
        t = t
        # (2) hidden conv layer
        t = F.relu(self.conv1(t))  # (28-5+0)/1+1=24 输入为b(batch_size)*1*28*28 输出为b*6*24*24 relu后shape不变
        t = F.max_pool2d(t, kernel_size=2, stride=2)  # (24-2+0)/2+1=12 输出为b*6*12*12
        t = self.bn1(t)
        
        # (3) hidden conv layer
        t = F.relu(self.conv2(t))  # (12-5+0)/1+1=8 输出为b*12*8*8 relu后shape不变
        t = F.max_pool2d(t, kernel_size=2, stride=2)  # (8-2+0)/2+1=4 输出为b*12*4*4
        
        # (4) hidden linear layer
        t = F.relu(self.fc1(t.reshape(-1, 12*4*4)))  # t.reshape后为b*192 全连接层后输出为b*120 relu后shape不变
        t = self.bn2(t)
        # (5) hidden linear layer
        t = F.relu(self.fc2(t))  # 全连接层后输出为b*60 relu后shape不变
        
        # (6) output layer
        t = self.out(t)  # 全连接层后输出为b*10 relu后shape不变
        return t

Transofrms
...

数据处理可以直接使用 torchvision.transforms 下的处理函数,包括均值,随机旋转,随机裁剪等等,也可以自己实现一些 pytorch 中没有实现的处理函数,下面拿一个分割网络的处理函数举例,可支持同时对传入的 Image 和 GroundTruth 进行处理,使用时直接按照顺序构造 ProcessImgAndGt 即可。

class ProcessImgAndGt(object):
    def __init__(self, transforms):
            self.transforms = transforms
    def __call__(self, img, label):
        for t in self.transforms:
            img, label = t(img, label)
        return img, label

class Resize(object):
    def __init__(self, height, width):
        self.height = height
        self.width = width
    def __call__(self, img, label):
        img = img.resize((self.width, self.height), Image.BILINEAR)
        label = label.resize((self.width, self.height), Image.NEAREST)
        return img, label

class Normalize(object):
    def __init__(self, mean, std):
        self.mean, self.std = mean, std
    def __call__(self, img, label):
        for i in range(3):
            img[:, :, i] -= float(self.mean[i])
        for i in range(3):
            img[:, :, i] /= float(self.std[i])
        return img, label

class ToTensor(object):
    def __init__(self):
        self.to_tensor = torchvision.transforms.ToTensor()
    def __call__(self, img, label):
        img, label = self.to_tensor(img), self.to_tensor(label).long()
        return img, label


transforms = ProcessImgAndGt([
    Resize(512, 512),
    Normalize([0.5, 0.5, 0.5], [0.1, 0.1, 0.1]),
    ToTensor()
])

Dataset
...

创建一个数据集类,继承 torch.utils.data.Dataset,只需实现init构造函数,getitem迭代器遍历函数以及len函数。

  • init函数中读取传入的数据集路径下的指定数据文件,还是拿一个分割网络的 dataset 流程举例,其他分类分类模型可以直接将 GroundTruth 替换为对应 label 即可,将拼接处理好的图片文件路径和 GroundTruth 文件路径作为元组存入一个为列表的成员变量 file_list 中;
  • getitem中根据传入的索引从 file_list 取对应的元素,并且通过 Transforms 进行处理;
  • len中返回 len(self.file_list) 即可。
class MyDataset(torch.utils.data.Dataset):
    def __init__(self, dataset_path, transforms):
        super(TrainDataset, self).__init()
        self.dataset_path = dataset_path
        self.transforms = transforms
        # 根据具体的业务逻辑读取全部数据路径作为加载数据的索引
        for dir in os.listdir(dataset_path):
            image_dir = os.path.join(dataset_path, dir)
            gt_path = image_dir + '/GT/'
            img_path = image_dir + '/Frame/'
            img_list = []
            for name in os.listdir(img_path):
                if name.endswith('.png'):
                    img_list.append(name)
            self.file_list.extend([(img_path + name, gt_path + name) for name in img_list])

    def __getitem__(self, idx):   
        img_path, label_path = self.file_list[idx]
        img = Image.open(img_path).convert('RGB')
        label = Image.open(label_path).convert('L')
        img, label = self.transforms(img, label)
        return img, label

    def __len__(self):
        return len(self.file_list)

Optimizer
...

选择优化器进行模型参数更新,要创建优化器必须给它一个可进行迭代优化的包含了全部参数的列表 然后可以指定针对这些参数的学习率(learning_rate),权重衰减(weight_decay),momentum 等,

optimizer = optim.Adam(model.parameters(), lr = 0.0001)

或者是可以指定针对哪些参数执行不一样的优化策略,根据不同层的 name 对不同层使用不同的优化策略。列表中的每一项都可以是一个 dict,dict 中 params 对应当前项的参数列表,可以对当前项指定学习率或者是衰减策略。对 base_params 使用的 1e-4 的学习率,对 finetune_params 使用 1e-3 的学习率,对两者一起使用 1e-4 的权重衰减

base_params = [params for name, params in model.named_parameters() if ("xxx" in name)]
finetune_params = [params for name, params in model.named_parameters() if ("yyy" in name)]
optimizer = optim.Adam([
    {"params": base_params},
    {"params": finetune_params, "lr": 1e-3}
], lr=1e-4, weight_decay=1e-4);

Run
...

基础组件都写好了,剩下的就是组成一个完整的模型结构。

  1. 实例化模型对象,并将其加载到 GPU 中
  2. 根据需要构建数据预处理对象,传入数据集对象中进行读取数据时的数据处理
  3. 构建训练和测试的数据集对象,并将其传入 torch.utils.data.DataLoader,指定 batch_size(训练或测试是每次读取多少条数据)、shuffle(读取数据时是否打乱)、num_workers(开启多少线程进行数据加载,为 0 时 (不推荐) 用主线程在训练模型时进行数据加载)等参数
  4. 使用 torch.optim.Adam 构建优化器对象,这里根据不同层的 name 对不同层使用不同的优化策略
  5. 训练 20 个 epoch,并且每 5 个 epoch 在测试集上跑一遍,这里只计算了损失,对于其他评价指标直接计算即可
  6. 根据条件对指定 epoch 的模型进行保存
  • optimizer.zero_grad() # pytorch 会积累梯度,在优化每个 batch 的权重的梯度之前将之前计算出的每个权重的梯度置 0
  • loss.backward() # 在最后一个张量上调用反向传播方法,在计算图中计算权重的梯度
  • optimizer.step() # 使用预先设置的学习率等参数根据当前梯度对权重进行更
  • model.train() # 保证 BN 层能够继续计算数据的均值和方差并进行更新,保证 dropout 层会按照设定的参数设置保留激活单元的概率(保留概率 = p)
  • model.eval() # BN 层会停止计算均值和方差,直接使用训练时的参数,dropout 层利用了训练好的全部网络连接,不随机舍弃激活单元
model = Network().cuda()
# 构建数据预处理
transforms = ProcessImgAndGt([
    Resize(512, 512),
    Normalize([0.5, 0.5, 0.5], [0.1, 0.1, 0.1]),
    ToTensor()
])
# 构建Dataset
train_dataset = MyDataset(train_dataset_path, transforms)
# DataLoader
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                                   batch_size=12,
                                                   shuffle=True,
                                                   num_workers=4,
                                                   pin_memory=False)
# TestDataset
test_dataset = MyDataset(test_dataset_path, transforms)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
                                                  batch_size=4,
                                                  shuffle=True,
                                                  num_workers=2,
                                                  pin_memory=False)
 # optimizer需要传入全部需要更新的参数名称,这里是对不用的参数执行不同的更新策略 
base_params = [params for name, params in model.named_parameters() if ("xxx" in name)]
finetune_params = [params for name, params in model.named_parameters() if ("yyy" in name)]
optimizer = torch.optim.Adam([
    {"params": base_params, "lr": 1e-3, ...},
    {"params": finetune_params, "lr": 1e-4, ...}
])

for epoch in range(20):
    model.train()
    epoch_loss = 0
    for batch in trian_loader:
        images. gts = batch[0].cuda(), batch[1].cuda()
        preds = model(iamges)
        loss = F.cross_entropy(preds, gts)
        optimizer.zero_grad()    # pytorch会积累梯度,在优化每个batch的权重的梯度之前将之前计算出的每个权重的梯度置0
        loss.backward()          # 在最后一个张量上调用反向传播方法,在计算图中计算权重的梯度 
        optimizer.step()         # 使用预先设置的学习率等参数根据当前梯度对权重进行更新
        epoch_loss += loss * trian_loader.batch_size
 # 计算其他标准
    loss = epoch_loss / len(train_loader.dataset)
 # .......
 # 每隔几个epoch在测试集上跑一下
    if epoch % 5 == 0:
        model.eval()
        test_epoch_loss = 0
        for test_batch in test_loader:
            test_images. test_gts = test_batch[0].cuda(), test_batch[1].cuda()
            test_preds = model(test_iamges)
            loss = F.cross_entropy(test_preds, test_gts)
            test_epoch_loss += loss * test_loader.batch_size
 # 计算其他标准
        test_loss = test_epoch_loss / (len(test_loader.dataset))
 # .......
 # 根据条件对指定epoch的模型进行保存 将模型序列化到磁盘的pickle包
    if 精度最高:
        torch.save(model.stat_dict(), f'{model_path}_{time_index}.pth')

Test
...

实际使用时需要将训练好的模型上在输入数据上运行,这里以测试集的数据为例,实际情况下只需要初始化模型之后将视频流中的图像帧作为模型的输入即可。

torch.no_grad()

  • 停止 autograd 模块的工作,不计算和储存梯度,一般在用训练好的模型跑测试集时使用,因为测试集时不需要计算梯度更不会更新梯度。使用后可以加速计算时间,节约 gpu 的显存
test_dataset = MyDataset(test_dataset_path, transforms)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
                                                       batch_size=1,
                                                       shuffle=False,
                                                       num_workers=2)
model = Network().cuda()
# 对磁盘上的pickle文件进行解包 将gpu训练的模型加载到cpu上
model.load_stat_dict(torch.load(model_path, map_location=torch.device('cpu')));
mocel.eval()

with torch.no_grad():
    for batch in test_loader:
        test_images. test_gts = test_batch[0].cuda(), test_batch[1].cuda()
        test_preds = model(test_iamges)
 # 保存模型输出的图片

石郎

pytorch 是非常好用的深度学习开源框架,使用 pytorch 编写深度学习模型,一般需要编写一下几个部分:

  • 模型定义
  • 数据处理和加载
  • 训练模型(Train and Validate)
  • 测试模型
  • 训练过程可视化(可选)

我以 pytorch 官网上的一份教程为例。这份教程是训练一个在数据集 cifar10 上的分类器。一共有十个类别分别为:‘airplane’, ‘automobile’, ‘bird’, ‘cat’, ‘deer’, ‘dog’, ‘frog’, ‘horse’, ‘ship’, ‘truck’。数据集中的图像大小为 33232。

  1. 模型的定义
import torch.nn as nn
import torch.nn.functional as F

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5) # 卷积
        self.pool = nn.MaxPool2d(2, 2) # 池化
        self.conv2 = nn.Conv2d(6, 16, 5) # 卷积
        # 全链接层,最后是输出10分类
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x
  1. 数据处理和加载

这一部分 pytorch 帮我们完成了,在torchvision.datasets中,已经实现了 ciar10 的数据集类。
如果我们需要自己定义的话,我们只需要继承torch.utils.Dataset类,重写其中的__getitem__方法和 __len__方法。

  1. 训练模型(Train and Validate)

训练模型一般分成以下几个步骤:

  • 定义网络

  • 定义数据

  • 定义损失函数和优化器

  • 开始训练

  • 训练网络

  • 将梯度置为 0

  • 求 loss

  • 反向传播

  • 更新参数

  • 更新优化器的学习率(可选)

  • 可视化各种指标

  • 计算在验证集上的指标 (可选)

# 定义网络
net = Net()

# 定义数据
#数据预处理,1.转为tensor,2.归一化
transform = transforms.Compose(    
     [transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
# 训练集
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,
                                          shuffle=True, num_workers=2)
# 验证集
testset = torchvision.datasets.CIFAR10(root='./data', train=False,
                                       download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=4,
                                         shuffle=False, num_workers=2)

# 定义损失函数和优化器 
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

# 开始训练
net.train()
for epoch in range(2):  # loop over the dataset multiple times
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        # get the inputs; data is a list of [inputs, labels]
        inputs, labels = data
        # 将梯度置为0
        # zero the parameter gradients
        optimizer.zero_grad()
        # 求loss
        # forward + backward + optimize
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        # 梯度反向传播
        loss.backward()
        # 由梯度,更新参数
        optimizer.step()

        # 可视化
        # print statistics
        running_loss += loss.item()
        if i % 2000 == 1999:    # print every 2000 mini-batches
            print('[%d, %5d] loss: %.3f' %
                  (epoch + 1, i + 1, running_loss / 2000))
            running_loss = 0.0

# 查看在验证集上的效果
dataiter = iter(testloader)
images, labels = dataiter.next()

# print images
imshow(torchvision.utils.make_grid(images))
print('GroundTruth: ', ' '.join('%5s' % classes[labels[j]] for j in range(4)))

net.eval()
outputs = net(images)
_, predicted = torch.max(outputs, 1)
print('Predicted: ', ' '.join('%5s' % classes[predicted[j]]
                              for j in range(4)))

一般经过这几个步骤一个网络就训练完成了。其他项目的流程也都基本和这个类似。

  1. 模型的参数

一个深度学习网络有很多的参数可以配置,一般分成以下三类:

  • 数据集参数(文件路径、batch_size 等)
  • 训练参数(学习率、训练 epoch 等)
  • 模型参数(输入的大小,输出的大小)

这些参数可以写一个类保存,也可以写一个字典,然后使用 json 保存,这些都是需要自己去实现的,但是这些都是一些细枝末节东西,写了几次,找到一个自己最喜欢的方式就可以,不是深度学习项目中必要的部分。

最后提供一个很好的 pytorch 学习的干货

谢谢!

Gemfield​

最接近解惑的项目应该是下面这个:

小朱

最近一年跟同学在写一个基于 PyTorch 的库,我在里面负责整个库的内核和接口设计,应该可以来回答这个问题。

首先深度学习的代码结构是没有一个统一标准的。一方面,代码结构取决于开发者自身的编程观念和水平,有人会一路长函数写到底,有人会利用面向对象进行封装和复用。另一方面,不同规模的项目,本身需要的结构也是很不一样的。Prototype 代码讲究简洁易懂,而平台级别的库讲究模块化和可维护性,这也是为什么大家看懂了 MNIST 上的代码,却经常看不懂开源库的原因。

就常见深度学习 project 来说,我觉得大概可以分成以下三类结构:

  1. Prototype 型:用最少的代码在 toy dataset 上实现一个模型。
  2. 可扩展型:围绕某个任务 / 某类模型展开的一套代码,和一些可选的模块。
  3. 平台库型:对一系列任务的一套统一代码。

Prototype 型的代码大家肯定都不陌生,比如 PyTorch 官方 tutorial 里的 CNNRNN,以及 Kipf 的 GCN。这类代码一般有效代码量在一两百行,主要分成三个部分:

  • 数据集读取和预处理
  • 模型定义
  • 模型训练、测试和保存代码

训练和测试代码一般写在主程序里,也有的会封装成叫 train 或者 test/inference 之类的函数。模型定义部分一般是一个独立文件,叫 model.py。数据集读取和预处理代码一般在 data.py 或者 utils.py 之类的文件里。

Prototype 型的代码的一大优点就是简单且好移植。当你想用另一个领域的技术时,比如做 vision 的用 GNN,拿一个这样 code 过来,是最容易吃透并且整合到自己 project 里的。当然 prototype 型的缺点也不少,一是缺少可选组件不易于刷点,二是魔改多了容易代码冗长混乱。

可扩展型一般是在 prototype 基础上,将数据集、模型单独抽象出来,把训练代码重构成模型无关的。这类代码多见于长期做同一个任务的开发 / 研究组,一般在千行左右,常见的模式是:

  • 各种数据集
  • 各种神经网络层 /loss/CUDA op
  • 各种 backbone 模型
  • 模型无关的训练、测试和保存代码
  • 主程序,解析命令行传递超参数

可扩展型代码是很多人做领域内的开发 / 研究的起点,比如折腾个新的神经网络层,或者部署到具体落地任务上。相比 prototype 型代码,这种代码更有魔改潜力,而且花一些时间的话也能吃透。

代码规模再往上翻个十倍,一般就是些平台库了。平台库的整体结构跟可扩展型并没有本质区别,但会多出很多细化的封装和测试模块,主要是因为大项目更需要考虑维护成本。记得有个经验定律是,一段混乱的代码的 debug 时间跟代码长度的平方成正比。所以平台库结构上的常见策略是把一段复杂代码拆成大量简单模块,把混乱的代码限制在每个模块内,同时用各种函数和类的抽象来避免重复代码。

这就解释了为什么一些看似很普通的代码,在平台里也会封装成特定的函数。比如题主给的 worker 函数,就是在多进程训练 / 测试里,单个进程的抽象的执行流程。对于只需要用库里某个模型,并且不用多卡训练的用户来说,这种设计的确是多余的。但从平台维护以及兼容不同用户需求的角度来看,却是必须的。类似的,config 文件主要是方便超参数搜索以及做 ablation study 用的。如果自己只是想开发一个小项目跑通一个模型,可以省去。

梦里风林

这问题简直就是为我这篇文章定制的

源码在这里: