是否可以在 pytorch 的 transform.compose 中添加自己的函数答案

【问题标题】：Is it possible to add own function in transform.compose in pytorch是否可以在 pytorch 的 transform.compose 中添加自己的函数
【发布时间】：2021-02-01 19:39:01
【问题描述】：

我正在使用预训练的 Alex 模型。我在一些随机图像数据集上运行这个模型。我想在训练之前将 RGB 图像转换为 YCbCr 图像。

我想知道是否可以自己给transform.compose添加一个函数，例如：

transform = transforms.Compose([
  ycbcr(), #something like this
  transforms.Resize((224, 224)),
  transforms.ToTensor(),
  transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

在哪里，

def ycbcr(img):
   img = cv2.imread(img)  
   img = cv2.cvtColor(img, cv2.COLOR_BGR2ycbcr)
   t = torch.from_numpy(img)
 return t

training_dataset = datasets.ImageFolder(link_train ,transform = transform_train)

training_loader = torch.utils.data.DataLoader(training_dataset, batch_size=96, shuffle=True)

这个过程正确吗？请帮助我如何进行？

【问题讨论】：

这里的transforms 是什么？ bgr2gray 也基本上是：gray = 0.2989*red + 0.5870*green + 0.1140*blue 所以你可以只使用矩阵乘法。
我很抱歉这个错误，我想在将其输入模型之前将所有图像更改为 YCbCr，因此我使用了变换。撰写，但它没有工作。我想在 transform.compose 中添加自己的函数。

标签： python-3.x machine-learning deep-learning neural-network pytorch

【解决方案1】：

您可以通过定义一个类将自定义转换传递给torchvision.transform。

为了更好地理解我建议你阅读documentations。

在您的情况下，它将类似于以下内容：

class ycbcr(object):
    def __call__(self, img):
        """
        :param img: (PIL): Image 

        :return: ycbr color space image (PIL)
        """
        img = cv2.imread(img)  
        img = cv2.cvtColor(img, cv2.COLOR_BGR2ycbcr)
        # t = torch.from_numpy(img)

        return Image.fromarray(t)

    def __repr__(self):
        return self.__class__.__name__+'()'

注意它获取了一个 PIL 图像并返回一个 PIL 图像。因此，您可能需要正确调整代码。但这是定义自定义转换的一般方式。

【讨论】：