torchvision.transforms GPU加速，提升预测阶段数据预处理速度

最新推荐文章于 2024-09-04 18:35:59 发布

胖胖大海

最新推荐文章于 2024-09-04 18:35:59 发布

阅读量7.3k

点赞数 13

分类专栏：深度学习 pytorch transforms 文章标签： pytorch加速 transform加速 torchvision torchvision GPU

本文链接：https://blog.csdn.net/cxx654/article/details/122610413

版权

深度学习同时被 3 个专栏收录

83 篇文章 5 订阅

订阅专栏

pytorch

16 篇文章 1 订阅

订阅专栏

transforms

1 篇文章 0 订阅

订阅专栏

参考：Tensor transforms and JIT — Torchvision 0.11.0 documentation

Torchvision v0.8.0之前版本：

Torchvision v0.8.0之前版本的transforms主要分为两类：

1、一类处理的输入数据类型为Tensor

2、另一类处理的数据类型为PILImage

所以torchvision里面提供了两个转换函数ToTensor和ToPILImage：

1、ToTensor：将PILImage（形状为 [H, W, C]）转换为Tensor（形状为[C, H, W]）

2、ToPILImage：将Tensor（形状为[C, H, W]）转换为PILImage（形状为 [H, W, C]）

这样转来转去的比较麻烦，并且这些转换操作基本上都在CPU上进行，处理速度比较慢。

Torchvision v0.8.0之后版本：

Torchvision v0.8.0之后的版本，将transforms，分为了三类：

1、有一部分只能处理Tensor（如Normalize）

2、一部分只能处理PILImage

3、还有一部分既能处理PILImage，也能处理Tensor

这样的话，我们可以将Tensor的处理部分迁移到GPU上去处理，大大提升数据的预处理速度。关于torchvision.transform更多的内容请参考：torchvision.transforms — Torchvision 0.10.0 documentation

Torchvision v0.8.0之前版本：

在模型预测阶段，我们通常按照模型的训练方式组织torchvision.transforms的pipeline，将pipeline的处理与训练阶段保持一致进行数据预处理，大致如下：

图上的transform pipeline比较简单，将读取的图片数据转换为PILImage，然后送入Resize，之后将PILImage转换为Tensor（注意，对于图片数据，ToTensor会除以255将数值缩放到[0,1]之间），最后normalize进行归一化，这一连串的操作全都在CPU上进行，实测发现ToTensor这步的耗时很长（上百毫秒），并且CPU占用2000%。

Torchvision v0.8.0之后版本：

在torchvision v0.8.0版本中新增了from torchvision.io import read_image函数用来读取图片文件，并将图片文件转换为Tensor对象（注意：与原来的transforms.ToTensor()不同，这里转换为Tensor并没有对图像数据缩放到[0, 1]之间），同时torchvision中的一些原本只能处理PILImage的transform在这个版本可以同时处理PILImage和Tensor了，如上面图片中的Resize。所以，我们现在就可以将原本对于PILImage的一些处理，转换为处理Tensor，并且将这些对于Tensor的处理放到GPU上来做。

对于上图中的old_transforms操作，转换为处理Tensor之后如下：

主要包括以下几步：

使用read_image读取文件，得到RGB图像数据的Tensor
将Tensor放到GPU上
在GPU上执行一系列能处理Tensor的transform

其中，ZeroOneNormalize是自定义的transform，这个transform用来实现将Tensor除以255转换到[0,1]之间，保持跟torchvision.transforms.ToTensor的处理一致，并且只有这样才能在下一步使用ImageNet预训练模型的Normalize归一化参数mean和std。ZeroOneNormalize转换算子也可以替换为torchvision.transforms.ConvertImageDtype，前提是上一步的输出类型是uint8。这样的话，这一连串的转换操作全在GPU上完成，速度提升嗖嗖的。

发现问题：

但是，实验验证发现，对于同一张图片，先转换为PILImage，然后进行Resize，与先转换为Tensor再进行Resize得到的结果之间存在一些差异，也就是得到的数值不同，这个在pytorch的官网上Resize函数的介绍部分也有说明。

备注：antialias参数在Torchvision v0.10.0才有，之前的版本都没有。

上图说，对于输入的PILImage和Tensor，Resize之后的结果会有轻微的不同，为了降低不同（并不能消除不同），提供了antialias参数，当输入为Tensor时，设置antialias=True，可以减缓这种不同程度。这种不同具体到什么程度，我测试了一张图片112x112大小，缩放到224x224之后，使用np.sum(img1 == img2)计算得到的相同的数值大概在11万，也就是说有大概4万个数值发生了变化，脆弱的模型感觉要崩。

所以WARNING部分也说了，建议在模型的训练和测试时使用相同的输入类型，也就是说建议在模型训练时就是用Tensor的形式，这样在测试阶段也是用Tensor，就保持跟训练时一致了。下面的代码示例就是将训练的transform也转换成处理Tensor的形式：