提示：文章寫(xiě)完后，目錄可以自動(dòng)生成，如何生成可參考右邊的幫助文檔

前言

最近在使用pytorch框架進(jìn)行模型訓(xùn)練時(shí)遇到一個(gè)性能問(wèn)題，即數(shù)據(jù)讀取的速度遠(yuǎn)遠(yuǎn)大于GPU訓(xùn)練的速度，導(dǎo)致整個(gè)訓(xùn)練流程中有大部分時(shí)間都在等待數(shù)據(jù)發(fā)送到GPU，在資源管理器中呈現(xiàn)出CUDA使用率周期性波動(dòng)，且大部分時(shí)間都是在等待數(shù)據(jù)加載。

一、造成的原因

其實(shí)從前言中就可以知道，造成這樣的原因可以理解為：GPU的算力與數(shù)據(jù)加載速度之間的不匹配。

二、查找不匹配的原因

本人使用的GPU為RTX3060，數(shù)據(jù)集為cifar10，使用的模型為VGG，顯然這張顯卡對(duì)于這個(gè)任務(wù)來(lái)說(shuō)是綽綽有余的，無(wú)論是顯存還是算力。
但是幾經(jīng)測(cè)試發(fā)現(xiàn)，數(shù)據(jù)從內(nèi)存送到GPU的速度實(shí)在是太慢了，去百度了很久都沒(méi)有很好的解決。那回到這個(gè)問(wèn)題的本身，既然是數(shù)據(jù)加載導(dǎo)致的性能差，那問(wèn)題就出在pytorch的dataset和dataloader中。

在dataset中，會(huì)將數(shù)據(jù)從磁盤(pán)讀入內(nèi)存中，如果啟用了dataloader中的pin_memory，就會(huì)讓數(shù)據(jù)常駐內(nèi)存，同時(shí)設(shè)置num_workers還能實(shí)現(xiàn)多進(jìn)程讀取數(shù)據(jù)，但即使設(shè)置了這些，數(shù)據(jù)加載速度依然沒(méi)有質(zhì)的提升。

博主發(fā)現(xiàn)，dataset中的transform是導(dǎo)致性能慢的一個(gè)原因，dataset中有個(gè)函數(shù)為_(kāi)_getitem__，每獲取一個(gè)數(shù)據(jù)就會(huì)讓這個(gè)數(shù)據(jù)過(guò)一次transform。相信大家都寫(xiě)過(guò)如下的代碼：

transform = transforms.Compose([
	transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize([0.4914, 0.4822, 0.4465], [0.5, 0.5, 0.5])
])

但是這里的ToTensor和Normalize完全沒(méi)必要沒(méi)讀一次數(shù)據(jù)都處理一次，可以在數(shù)據(jù)加載到內(nèi)存的時(shí)候就直接全部處理完，這樣每個(gè)數(shù)據(jù)只需要經(jīng)歷一次ToTensor和Normalize，這會(huì)大大提高數(shù)據(jù)讀取速度，大家可以自行測(cè)試一次ToTensor和Normalize所需要的時(shí)間，還是非常大的。

在訓(xùn)練的過(guò)程中，相信大家也寫(xiě)過(guò)如下代碼：

for x, y in dataloader:
	x, y = x.cuda(), y.cuda()

經(jīng)過(guò)博主測(cè)試，將數(shù)據(jù)發(fā)送到GPU也是非常耗時(shí)的，那為什么不一次性全部加載到GPU里面呢？當(dāng)然前提是你的GPU顯存夠大。

三、解決方法

以上分析可以總結(jié)為兩點(diǎn)：

1. 數(shù)據(jù)的預(yù)處理有一部分可以提前對(duì)全部數(shù)據(jù)做一遍；
2. 如果顯存足夠，可以將數(shù)據(jù)全部加載到GPU中。

基于此，我們可以重載類(lèi)，這里以pytorch自帶的cifar10為例：

class CUDACIFAR10(CIFAR10):
    def __init__(
            self, 
            root: str, 
            train: bool = True,
            to_cuda: bool = True,
            half: bool = False,
            pre_transform: Optional[Callable] = None,
            transform: Optional[Callable] = None, 
            target_transform: Optional[Callable] = None, 
            download: bool = False) -> None:

        super().__init__(root, train, transform, target_transform, download)
        if pre_transform is not None:
            self.data = self.data.astype("float32")
            for index in range(len(self)):
                """
                ToTensor的操作會(huì)檢查數(shù)據(jù)類(lèi)型是否為uint8, 如果是, 則除以255進(jìn)行歸一化, 這里data提前轉(zhuǎn)為float,
                所以手動(dòng)除以255.
                """
                self.data[index] = pre_transform(self.data[index]/255.0).numpy().transpose((1, 2, 0))
                self.targets[index] = torch.Tensor([self.targets[index]]).squeeze_().long()
                if to_cuda:
                    self.targets[index] = self.targets[index].cuda()
            self.data = torch.Tensor(self.data).permute((0, 3, 1, 2))
            if half:
                self.data = self.data.half()
            if to_cuda:
                self.data = self.data.cuda()

    def __getitem__(self, index: int) -> Tuple[Any, Any]:
        """
        Args:
            index (int): Index

        Returns:
            tuple: (image, target) where target is index of the target class.
        """
        img, target = self.data[index], self.targets[index]

        if self.transform is not None:
            img = self.transform(img)
        if self.target_transform is not None:
            target = self.target_transform(target)
        return img, target

to_cuda為T(mén)rue就會(huì)讓數(shù)據(jù)全部加載到GPU中，后續(xù)就不需要寫(xiě)x, y = x.cuda(), y.cuda()了。
pre_transform就是讓所有數(shù)據(jù)提前進(jìn)行的處理，例如使用ToTensor和Normalize，后續(xù)調(diào)用時(shí)不會(huì)再進(jìn)行這些處理。
transform為后續(xù)調(diào)用時(shí)會(huì)進(jìn)行的處理，一般就是一些隨機(jī)處理過(guò)程。

在博主的測(cè)試過(guò)程中發(fā)現(xiàn)，解決了以上問(wèn)題后，一個(gè)epoch只要2秒就能完成，而平時(shí)需要15秒，并且任務(wù)管理器中的CUDA幾乎全程拉滿。唯一的代價(jià)就是顯存占用更高了，這何嘗不是一種空間換時(shí)間。

四、使用方法

這里直接粘貼我為這個(gè)類(lèi)寫(xiě)的注釋

- 使用pytorch自帶的CIFAR10時(shí), 每讀取一個(gè)數(shù)據(jù)都會(huì)調(diào)用一次transforms, 其中ToTensor()和Normalize()會(huì)消耗巨大的時(shí)間
  如果你的數(shù)據(jù)集非常的大, 那么一個(gè)epoch將會(huì)花費(fèi)非常多的時(shí)間用于讀取數(shù)據(jù), 如果還要將數(shù)據(jù)送入GPU, 那么時(shí)間將會(huì)繼續(xù)增加。

- 一般的寫(xiě)法如下：

    for epoch in range(epochs):
        for x, y in dataloader:
            x, y = x.cuda(), y.cuda()

  如果你的數(shù)據(jù)集很大并且GPU算力很強(qiáng), 那么讀取數(shù)據(jù)并發(fā)送的GPU將會(huì)成為性能瓶頸。

- CUDACIFAR10是專(zhuān)門(mén)針對(duì)pytorch的CIFAR10進(jìn)行優(yōu)化的, 使用的前提是你的顯存足夠的大, 至少8GB, 且讀取數(shù)據(jù)已經(jīng)是性能瓶頸。
  CUDACIFAR10的參數(shù)與CIFAR10非常相似, 新增的參數(shù)為：
    to_cuda: bool, 是否將數(shù)據(jù)集常駐GPU, default=True
    half: bool, 進(jìn)一步降低數(shù)據(jù)所占據(jù)的顯存, 在混合精度訓(xùn)練時(shí)使用, 否則可能存在意外(例如梯度值overflow)
    pre_transform: 傳入一個(gè)transforms, 如果不為None, 則會(huì)在初始化數(shù)據(jù)時(shí)直接對(duì)所有數(shù)據(jù)進(jìn)行對(duì)應(yīng)的轉(zhuǎn)換, 在后續(xù)調(diào)用時(shí)將
                    不會(huì)使用pre_transform. 可以將ToTensor()和Normalize()作為pre_transform, 這會(huì)大幅度減少讀取時(shí)間.

- CUDACIFAR10的用法如下:

    pre_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize([0.4914, 0.4822, 0.4465], [0.5, 0.5, 0.5])
    ])
    dataset = CUDACIFAR10(..., to_cuda=True, pre_transform=pre_transform, ...)
    dataloader = Dataloader(dataset, ..., pin_memory=False, num_workers=0, ...)
    ...
    for epoch in range(epochs):
        for x, y in dataloader:
            # 不需要寫(xiě)x, y = x.cuda(), y.cuda(), 除非to_cuda=False
            ...

- 使用CUDACIFAR10需要注意如果啟用了to_cuda, 那么Dataloader不能啟用pin_memory, pin_memory是將數(shù)據(jù)常駐內(nèi)存, 這會(huì)產(chǎn)生沖突.
  同時(shí)num_workers=0.

- 如果參數(shù)to_cuda=False, pre_transform=None, 那么該類(lèi)與CIFAR10用法完全一致.

后言

本文寫(xiě)作倉(cāng)促，可能有部分錯(cuò)誤，歡迎大家的批評(píng)與指正。文章來(lái)源地址http://www.zghlxwxcb.cn/news/detail-410887.html

到了這里，關(guān)于解決pytorch中Dataloader讀取數(shù)據(jù)太慢的問(wèn)題的文章就介紹完了。如果您還想了解更多內(nèi)容，請(qǐng)?jiān)谟疑辖撬阉鱐OY模板網(wǎng)以前的文章或繼續(xù)瀏覽下面的相關(guān)文章，希望大家以后多多支持TOY模板網(wǎng)！