Pytorch 加载单通道图片用作数据集训练和测试

hfw6310

于 2020-06-28 16:42:11 发布

阅读量1.6k

点赞数 1

CC 4.0 BY-SA版权

本文链接：https://blog.youkuaiyun.com/hfw6310/article/details/107002787

本文介绍如何使用PyTorch中的transforms.Grayscale(1)将彩色图像转换为灰度图像，并展示了如何在训练和验证数据加载器中应用此转换。通过调整transforms.Compose中的参数，可以实现对图像数据集的预处理，包括尺寸调整、水平翻转和张量转换。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

用transforms.Grayscale(1)来实现。如下所示：
train_dataset = datasets.ImageFolder(
    traindir,
    transforms.Compose([
        #transforms.RandomResizedCrop(224),
        transforms.Resize((112,112)),
        transforms.RandomHorizontalFlip(),
        transforms.Grayscale(1),  ############
        transforms.ToTensor(),
    ]))

train_sampler = None

train_loader = torch.utils.data.DataLoader(
    train_dataset, batch_size=args.batch_size, shuffle=True,
    num_workers=args.workers, pin_memory=True, sampler=train_sampler)

val_loader = torch.utils.data.DataLoader(
    datasets.ImageFolder(valdir, transforms.Compose([
        transforms.Resize((112,112)),
        #transforms.CenterCrop(224),
        transforms.Grayscale(1), ############
        transforms.ToTensor(),
    ])),
    batch_size=args.batch_size, shuffle=False,
    num_workers=args.workers, pin_memory=True)