【Pytorch】conv2d torchvision.transforms

• 示例代码

from torchvision import transforms
from torch.utils.data import Dataset, DataLoader
from PIL import Image
from torch import nn
import os


class ImgDataset(Dataset):
    def __init__(self, root_dir):
        self.root_dir = root_dir
        self.img_path_list = os.listdir(root_dir)
        # 统一图片维度，转为tensor格式
        self.compose = transforms.Compose([transforms.Resize((500, 500)), transforms.ToTensor()])

    def __getitem__(self, idx):
        img_path = os.path.join(self.root_dir, self.img_path_list[idx])
        img = Image.open(img_path)
        # img = transforms.ToTensor()(img) # 图片维度不一致，dataloader不能遍历
        img = self.compose(img)
        return img

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


class MyConv(nn.Module):
    def __init__(self):
        super(MyConv, self).__init__()
        self.conv2d = nn.Conv2d(in_channels=3, out_channels=6, kernel_size=3, stride=1)

    def forward(self, data):
        data = self.conv2d(data)
        return data


dataset = ImgDataset('./img_data')
dataloader = DataLoader(dataset=dataset, batch_size=2, shuffle=False, drop_last=False)
test_conv2d = MyConv()

print('befor conv', '=' * 20)
for i in dataloader:
    print(i.shape)

print('after conv', '=' * 20)
for i in dataloader:
    print(test_conv2d(i).shape)
    

• img_data