【pytorch】给dataloader提速：Kornia库 + GPU

简介

目前在做一个视频复原的相关工作，使用pytorch构建深度学习模型。在训练的过程中发现，数据集总共有6、7千对，使用单张2080TI训练600epoch至少要大概5天时间，而且时间很不稳定，有的训练可能要9天时间。这么长的训练周期让我不得不放几天假。

在debug的时候发现，数据加载和预处理的过程花费了大量时间，因为计算都是在CPU上完成的。此外，预处理使用的multi progress经常会相互等待，导致时间很不稳定，而且同时跑多个代码可能会导致死锁。因此我想着使用GPU来加速。

GPU加速当然使用别人写好的库啦！我最开始找的库是NVIDIA的DALI，但是这个库比较新，参考资料比较少，官方文档也没马上看懂，问师兄师姐也没人用过。因此就先转用另一个库Kornia——pytorch版的opencv。

简单来说，我的提速方案就是用定义model的方式定义了一个GPU的transform，图片通过网络的fooward()之前，先通过transform的forward()

原始dataloader

原始dataload类继承了torch.utils.data.dataset.Dataset类，使用opencv来读取图片。train阶段的预处理操作按照顺序为 crop → color jitter → normalize → flip (vertical & horizontal) → color channel change → gaussian noise → to tensor ，test阶段为。同一视频序列的图片采用的预处理操作都相同。所有操作都在CPU上完成。示意代码如下：

class VideoDataset(torch.utils.data.dataset.Dataset):
    def __init__(self, imags_path, transform=None):
        self.transform = transform
        ....
    def __getitem__(self, ...):
        imgs = []
        for path in self.imags_path:
            img = cv2.imread(path ).astype(np.float32)
            imgs.append(self.transforms(img))
        return imgs		# 返回的是图片序列
        
if __name__ == '__main__':
    train_transforms = Compose([RandomCrop(cfg.DATA.CROP_IMG_SIZE),
                                ColorJitter(cfg.DATA.COLOR_JITTER),
                                Normalize(mean=cfg.DATA.MEAN,std=cfg.DATA.STD),
                                RandomVerticalFlip(),
                                RandomHorizontalFlip(),
                                RandomColorChannel(),
                                RandomGaussianNoise(cfg.DATA.GAUSSIAN),
                                ToTensor()])
    test_transforms = Compose([Normalize(mean=cfg.DATA.MEAN, std=cfg.DATA.STD),
                               ToTensor()])
    train_data_loader = torch.utils.data.DataLoader(dataset=VideoDataset(imags_path, train_transforms),
                                                    batch_size=1, num_workers=2, pin_memory=True, shuffle=True)
    test_data_loader = torch.utils.data.DataLoader(dataset=VideoDataset(imags_path, test_transforms),
                                                    batch_size=1, num_workers=2, pin_memory=True, shuffle=True)
    for i in range(epoch):
        for idx, imgs in enumerate(train_data_loader):
            for img in imgs：
            	img = img.cuda()
                # 然后就是前向传播、算loss、反向传播。。。
                ...

提速后的dataloader

我仍然使用一个继承了torch.utils.data.dataset.Dataset的dataload类。其中，使用opencv读取图片，然后crop，最后转换成tensor。示意代码如下

class VideoDataset(torch.utils.data.dataset.Dataset):
    def __init__(self, ...):
        ....
    def __getitem__(self, ...):
        img = cv2.imread(img_path).transpose(2,0,1)
        img = img[:, crop_y1: crop_y2 + 1, crop_x1: crop_x2 + 1] # crop
        return torch.from_numpy(img_blur).float()				 # to tensor

其他的预处理操作我则放在训练阶段完成。首先，我使用Kornia库定义了一个继承了nn.Module的Transforms类，里面完成如下操作color jitter → flip (vertical & horizontal) → color channel change → gaussian noise。相当于用构建网络的方式构建预处理的操作。示意代码如下

import kornia as K
class Transforms(nn.Module):
    def __init__(self):
        ...
    def forward(self, img):
        # color jitter
        img = K.enhance.adjust_brightness(img, self.brightness)
        img = K.enhance.adjust_contrast(img, self.contrast)
        img = self.adjust_saturation_hue(img, self.saturation, self.hue)
        # flip
        img = K.geometry.transform.hflip(img)
        img = K.geometry.transform.vflip(img)
        # Channel Shuffle
        img = img[:, self.channel_shuffle_order, :, :]
        # gaussian noise
        img = img + self.gaussian_noise
        return img

然后，在遍历数据集拿到每一张图片后，首先转到GPU，再通过Transforms即可。示意代码如下

if __name__ == '__main__':
    train_data_loader = torch.utils.data.DataLoader(dataset=VideoDataset(...), batch_size=1, 
                                                    num_workers=2, pin_memory=True, shuffle=True)
    for i in range(epoch):
    	for idx, imgs in enumerate(train_data_loader):
    		train_transform = Transforms().cuda()
            for img in imgs：
            	img = train_transform(img.cuda())
            	# 然后就是前向传播、算loss、反向传播。。。
                ...

效果

我使用数据集测试了一下提速前后遍历数据集的耗时。测试方法就是将正常训练model的代码去掉前向传播、计算loss、反向传播等操作，只保留数据的加载、预处理、转移到GPU的操作。是用的数据集中总共有1万张图片。

原始的dataload在10epoch下总耗时11904s（下图一），加速后的dataload在10epoch下耗时791s（下图二）。此外，可以看到原始的dataload各个epcoh的耗时很不稳定，短的能有150s，长的能有4000s；而加速后的dataload耗时基本上都在80s左右。

完整代码

前面的code都是示意代码，原始和提速后的dataloader的完整代码我已经上传到Github上。

原始dataloader来自论文STFAN，我测试用的数据集是DVD数据集

其他尝试

我原本是想在VideoDataset类中将每一张图片转移到GPU上，然后是用Kornia库完成预处理。但是会报错