背景与需求

现在做深度学习的越来越多人都有用PyTorch，他容易上手，而且API相对TF友好的不要太多。今天就给大家带来最近PyTorch训练的一些小小的心得。

大家做机器学习、深度学习都恨不得机器卡越多越好，这样可以跑得越快，道理好像也很直白，大家都懂。实际上我们在训练的时候很大一部分制约我们的训练的速度快慢被IO限制住了，然面CPU的利用率却不高，就算有8卡了，然而GPU的利用率却长期处理低水平，不能发挥设备本应该有的水平。所以我一直在想，有什么办法能加快IO的读取，当然最直截的就换SSD，那上速度会直接上去了。那如果是我们在服务器或者是普通的电脑就没有办法呢吗？

而且经常用PyTorch的人应该会发现，如果我们把DataLoader的num_workers设置比较大的时候，在训练启动时会等待比较久，而且在每一个epoch之间的切换也是需要等挺久的（更换，加载数据）。

如果是一个程序员的话，肯定会想到多线程、多进程，这是否会能加速我们训练的IO？答案是肯定的。

今天给大家带来的就是，多线程读取数据的实例，本次测试不含训练部分，只是对Dataset, DataLoader数据加载的部分进行测试。

PyTorch DataLoader会产生一个index然后Dataset再进行读取，如果一个batch_size=128的话，那就要产生128次的数据调试，并读取。

我的想法就很简单，我想要不我就直接在Dataset就生成好所需的Batches，这样在DataLoader的batch_size=1的话，那也是对应一个batch的数据，而我在Dataset的可以用线程去加载数据，这样应该能提高读取的效率。

有了想法就是干了。

平时我们重要Dataset的结构如下，这里用到了albumentations作为数据处理的库，而不是torchvision的transforms，其它没有什么区别的

def default_loader(path):
    return Image.open(path).convert('RGB')

class AlbumentationsDatasetList(Dataset):
    """ 
        Data processing using albumentation same as torchvision transforms
    """
    def __init__(self, imgs, transform=None, loader=default_loader, percentage=1):
        # here can control the dataset size percentage    
        img_num = int(len(imgs) * percentage)
        self.imgs = imgs[:img_num]
        self.transform = transform
        self.loader = loader
        
    def __getitem__(self, index):
        fn = self.imgs[index]
        img = self.loader(fn)

        if self.transform is not None:
            image_np = np.array(img)
            augmented = self.transform(image=image_np)
            img = augmented['image']
        return img

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

  

 

  
1
  
2
  
3
  
4
  
5
  
6
  
7
  
8
  
9
  
10
  
11
  
12
  
13
  
14
  
15
  
16
  
17
  
18
  
19
  
20
  
21
  
22
  
23
  
24
  
25
  
26
 

方法的实现

说干就干，把多线程加进来进行改造Dataset，下面来看一下代码，代码加入了一些细节，所以会比较长，但结构还是跟上面的是一样的。只是Dataset就已经把batches都处理好了，在加载数据后，是把他们都stack在一起，这样就可以形成[N, C, W, H]结构的数据了。

注意：如果drop_last=False的话，那么最后的一个batch的数量一般不会与batch_size相同，所以在DataLoader的里batch_size要设置成1。还有DataLoader设置成1后，实际加载的数据是[1, N, C, W, H]，所以在用的时候要squeeze一下。

class AlbumentationsDatasetList(Dataset):

    def __init__(self,
                 images,
                 batch_num,
                 percentage=1,
                 transform=None,
                 multi_load=True,
                 shuffle=True,
                 seed=None,
                 drop_last=False,
                 num_workers=4,
                 loader=default_loader) -> None:

        # ==============================================
        # Set seed
        # ==============================================
        if seed is None:
            self.seed = np.random.randint(0, 1e-6, 1)
        else:
            self.seed = seed
        random.seed(self.seed)

        self.images = images
        self.batch_num = batch_num  # use batch_num instead of batch_size, same thing
        self.percentage = percentage
        self.transform = transform
        self.multi_load = multi_load
        self.shuffle = shuffle
        self.drop_last = drop_last
        self.num_workers = num_workers  # Dataset num_workers
        self.loader = loader

        self.batches = self._create_batches()
        self.batches = self._get_len_batches(self.percentage)

    def _get_len_batches(self, percentage):
        """
        Description:
            - you could control how many batches you want to use for training or validating
              indices sort, so that could keep the batches got in order from originla batches

        Parameters:
            - percentage: float, range [0, 1]

        Return
            - numpy array of the new bags
        """
        batch_num = int(len(self.batches) * percentage)
        indices = random.sample(list(range(len(self.batches))), batch_num)
        indices.sort()
        new_batches = np.array(self.batches, dtype='object')[indices]
        return new_batches

    def _create_batches(self, ):
        if self.shuffle:
            random.shuffle(self.images)

        batches = []
        ranges = list(range(0, len(self.images), self.batch_num))
        for i in ranges[:-1]:
            batch = self.images[i:i + self.batch_num]
            batches.append(batch)

        # == Drop last ===============================================
        last_batch = self.images[ranges[-1]:]
        if len(last_batch) == self.batch_num:
            batches.append(last_batch)
        elif self.drop_last:
            pass
        else:
            batches.append(last_batch)

        return batches

    def __getitem__(self, index):

        batch = self.batches[index]

        # == Stack all images, become a 4 dimensional tensor ===============
        if self.multi_load:
            batch_images = self._multi_loader(batch)

        else:
            batch_images = []
            for image in batch:
                img = self._load_transform(image)
                batch_images.append(img)

        batch_images_tensor = torch.stack(batch_images, dim=0)

        return batch_images_tensor

    def _load_transform(self, tile):
        img = self.loader(tile)

        if self.transform is not None:
            image_np = np.array(img)
            augmented = self.transform(image=image_np)
            img = augmented['image']
        return img

    def _multi_loader(self, tiles):
        images = []
        executor = ThreadPoolExecutor(max_workers=self.num_workers)
        results = executor.map(self._load_transform, tiles)
        executor.shutdown()
        for result in results:
            images.append(result)
        return images

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


  

 

  
1
  
2
  
3
  
4
  
5
  
6
  
7
  
8
  
9
  
10
  
11
  
12
  
13
  
14
  
15
  
16
  
17
  
18
  
19
  
20
  
21
  
22
  
23
  
24
  
25
  
26
  
27
  
28
  
29
  
30
  
31
  
32
  
33
  
34
  
35
  
36
  
37
  
38
  
39
  
40
  
41
  
42
  
43
  
44
  
45
  
46
  
47
  
48
  
49
  
50
  
51
  
52
  
53
  
54
  
55
  
56
  
57
  
58
  
59
  
60
  
61
  
62
  
63
  
64
  
65
  
66
  
67
  
68
  
69
  
70
  
71
  
72
  
73
  
74
  
75
  
76
  
77
  
78
  
79
  
80
  
81
  
82
  
83
  
84
  
85
  
86
  
87
  
88
  
89
  
90
  
91
  
92
  
93
  
94
  
95
  
96
  
97
  
98
  
99
  
100
  
101
  
102
  
103
  
104
  
105
  
106
  
107
  
108
  
109
  
110
  
111
  
112
  
113
  
114
 

代码与数据测试

接下来就是拿数据进行测试了，这里还设置了multi_load的参数，这样我们可以方便控制是否用多线程与否，这样我们就可以对比一下在相同的机器，相同的数据下，多线程加载数据是否比单线程快。

测试的目的：

1，是否多线程多单线程快；
2，多线程能比单线路快多少；
3，找到这机器最快（或者比较全适）的越参数，可作为其它机器的参考。
测试平台：Window10

CPU：Intel Core i7-9850H @ 2.60GHz

RAM: 32 GB

测试的数据：是5000张图像，全部都是3通道RBG，8位的512x512像素图像，图像格式是.PNG。

测试方法：

超参数如下：搜索空间为1024

multi_loads = [True, False]
prefetch_factors = list(range(0, 17, 2))[1:] # [2, 4, 6, 8, 10, 12, 14, 16]
dataset_workers = list(range(0, 17, 2))[1:]
dataloader_workers = list(range(0, 17, 2))[1:]

  

 

  
1
  
2
  
3
  
4
 

利用grid search方法，每一个搜索空间都对Dataset, DataLoader设置不同的参数，而且每轮数据都是读完、并处理完5000张图像，drop_last=False

数据增强：只做了resize，normalize

下面是全部的测试代码。

albumentations_valid = album.Compose([
        album.Resize(480, 480),
        album.Normalize(mean=[0.7347, 0.4894, 0.6820, ], std=[0.1747, 0.2223, 0.1535, ]),
        ToTensorV2(),
    ])

    from utils import get_specified_files
    path = r"xxxxx"
    images = get_specified_files(path, suffixes=[".png"], recursive=True) # glob.glob
    images = images[:5000]
    print(len(images))

    results = []
    
    log_file = open(r"grid_search_log.txt", mode='a', encoding='utf-8')

    multi_loads = [True, False]
    prefetch_factors = list(range(0, 17, 2))[1:] # [2, 4, 6, 8, 10, 12, 14, 16]
    dataset_workers = list(range(0, 17, 2))[1:]
    dataloader_workers = list(range(0, 17, 2))[1:]

    for multi_load in multi_loads:
        for prefetch_factor in prefetch_factors:
            for dataset_worker in dataset_workers:
                for dataloader_worker in dataloader_workers:

                    multi_load = multi_load
                    if multi_load:
                        prefetch_factor = prefetch_factor
                    else:
                        prefetch_factor = prefetch_factor

                    dataloader_worker = dataloader_worker

                    train_dataset = AEDataset(images, 
                                            batch_num=128, 
                                            percentage=1, 
                                            transform=albumentations_valid, 
                                            multi_load=multi_load, 
                                            shuffle=True, 
                                            seed=0, 
                                            drop_last=False,
                                            num_workers=dataset_worker,
                                            )
                    train_loader = DataLoader(dataset=train_dataset, 
                                            batch_size=1, 
                                            shuffle=False, 
                                            num_workers=dataloader_worker, 
                                            pin_memory=True, 
                                            prefetch_factor=prefetch_factor, 
                                            persistent_workers=False)
                    print("Start loading")
                    start_time = time.time()
                    for i, (batches) in enumerate(train_loader):
                        i+1
                    elapse = time.time() - start_time
                    
                    print(f"multi_load: {multi_load}, prefetch_factors: {prefetch_factor}, dataset_workers: {dataset_worker}, data_loader_workers: {dataloader_worker}, elapse: {elapse:.4f}")
                    log_file.write(f"multi_load: {multi_load}, prefetch_factors: {prefetch_factor}, dataset_workers: {dataset_worker}, data_loader_workers: {dataloader_worker}, elapse: {elapse:.4f}\n")

  

 

  
1
  
2
  
3
  
4
  
5
  
6
  
7
  
8
  
9
  
10
  
11
  
12
  
13
  
14
  
15
  
16
  
17
  
18
  
19
  
20
  
21
  
22
  
23
  
24
  
25
  
26
  
27
  
28
  
29
  
30
  
31
  
32
  
33
  
34
  
35
  
36
  
37
  
38
  
39
  
40
  
41
  
42
  
43
  
44
  
45
  
46
  
47
  
48
  
49
  
50
  
51
  
52
  
53
  
54
  
55
  
56
  
57
  
58
  
59
 

测试结果

回到我们上面的测试目标

测试的目的：

1，是否多线程多单线程快；
2，多线程能比单线程快多少；
3，找到这台机器最快（或者比较全适）的越参数，可作为其它机器的参考。
我们带着这3个问题，看一下下面的测试结果：

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

  

 

  
1
  
2
  
3
 

path = "C:/Users/jasne/Desktop/grid_search_multi_load.csv"
df = pd.read_csv(path)
df.head()

  

 

  
1
  
2
  
3
 

Max elapse

也是我们平时用的普通load的方法，时间是72.28秒

df[df["elapse"]==df["elapse"].max()]

  

 

  
1
 

Multi Load Max elapse

多线程时最慢的时间

multi_load = df[df["multi_load"]==True]
multi_load[multi_load["elapse"]==multi_load["elapse"].max()]

  

 

  
1
  
2
 

Min elapse

相差的倍数的计算公式为( max − min ) / min (\text{max} - \text{min}) / \text{min}(max−min)/min
时间是19.97秒，比最长的时间少了 52.31秒，快了2.6倍的时间，所以可以看出用multi_load肯定是比single load要快的。

多线程的时间，也受prefetch_factors, dataset_workers, dataloader_workers的影响。而且影响还是比较大的。

多线程时，最快与最慢的相差1.42倍

df[df["elapse"]==df["elapse"].min()]

  

 

  
1
 

下面来看是否 data_loader_workers越大越好？

dataloader_workers = multi_load[(multi_load["prefetch_factors"]==2) & (multi_load["dataset_workers"]==2)]
dataloader_workers.sort_values("data_loader_workers", inplace=True)
dataloader_workers

  

 

  
1
  
2
  
3
 

plt.figure(figsize=(8, 5))
plt.scatter(dataloader_workers["data_loader_workers"], dataloader_workers["elapse"])
plt.show()

  

 

  
1
  
2
  
3
 

从图上可以看出，dataloader_workers并非越大越好，dataloader_workers=4时是在2-8之间是比较好的选择。随着dataloader_workers的增加，所需要的时间也呈线性的增加。

下面来看是否 dataset_workers越大越好

dataset_workers = multi_load[(multi_load["prefetch_factors"]==2) & (multi_load["data_loader_workers"]==2)]
dataset_workers.sort_values("dataset_workers", inplace=True)
dataset_workers

  

 

  
1
  
2
  
3
 

plt.figure(figsize=(8, 5))
plt.scatter(dataset_workers["dataset_workers"], dataset_workers["elapse"])
plt.show()

  

 

  
1
  
2
  
3
 

从图上可以看出，dataset_workers增加也可以明显减少数据加载所需要时间。但是当dataset_workers超过10后，不再呈现出减少的趋势，当达到12、14时有一点点上降。由于测试平台有限，这里所应该让测试一下dataset_workers达到128或者更高的数之间，是否会达到更少的数据加载时间。

下面来看是否 prefetch_factors越大越好

prefetch_factors = multi_load[(multi_load["dataset_workers"]==2) & (multi_load["data_loader_workers"]==2)]
prefetch_factors.sort_values("prefetch_factors", inplace=True)
prefetch_factors

  

 

  
1
  
2
  
3
 

plt.figure(figsize=(8, 5))
plt.scatter(prefetch_factors["prefetch_factors"], prefetch_factors["elapse"])
plt.show()

  

 

  
1
  
2
  
3
 

从图上可以看出，prefetch_factors似乎好像越大，加载的时间越少，但似乎也相差不多，最多的时间与最小的时间相差也仅为2.6秒。

prefetch_factors的外一个筛选条件
prefetch_factors = multi_load[(multi_load["dataset_workers"]==10) & (multi_load["data_loader_workers"]==4)]
prefetch_factors.sort_values("prefetch_factors", inplace=True)
prefetch_factors

  

 

  
1
  
2
  
3
  
4
 

plt.figure(figsize=(8, 5))
plt.scatter(prefetch_factors["prefetch_factors"], prefetch_factors["elapse"])
plt.show()

  

 

  
1
  
2
  
3
 

从图上可以看出，prefetch_factors数量似乎对加载时间的影响似乎不太明显，最多的时间与最小的时间相差也仅为2.6秒。

plt.figure(figsize=(8, 5))
plt.scatter(prefetch_factors["prefetch_factors"], 
prefetch_factors["elapse"])
plt.show()

  

 

  
1
  
2
  
3
  
4
 

从图上可以看出，prefetch_factors数量似乎对加载时间的影响似乎不太明显，最多的时间与最小的时间相差也仅为2.6秒。

结论

多线程加载数据肯定是比单线程快的？
这点是不用质疑的，单从计算机的运行方式就可以得出这个结论，这也是并行的优势。
多线程能比单线程快多少？
从上面的结果，我们看到，当选用合适的超参数时，多线程加载相同的数据与相同的处理方法，比单线程快了52.31秒，快了2.6倍有多。就算是最不好的参数，多线和最长的加载时间为48.33秒，也比单线程的72.28秒，快差不多0.5倍。
找到这台机器最快（或者比较全适）的越参数，可作为其它机器的参考
dataset_workers 越大越好，但达到了一个临界值后，不会再增加了，本测试平台的值为10
data_loader_workers，不是越大越好，本测试平台最好的值为4，在4左右的值都是较好的参考值。然后随着此参数的数量的增加，所需要的时间也呈线性的增涨，这也说明了PyTorch大data_loader_workers启动需要等待更久的时间
prefetch_factors的数量似乎对数据的加载时间影响不大，但最好不要是1。
本次测试没有监测内存还有CPU的使用率，但在过程中观察了一下，CPU使用率基本都可以达到100%。也可以把这些参数也监测起来，形成更多的超参数，以便参考。
注意：由于在训练的过程中也是需要利用CPU的，所以尽量不要太多的dataset_workers，尽量不要把CPU都使用到100%，而造成死机。

文章来源: wanghao.blog.csdn.net，作者：AI浩，版权归原作者所有，如需转载，请联系作者。

原文链接：wanghao.blog.csdn.net/article/details/120810103