PyTorch实现重写/改写Dataset并载入Dataloader

下面是PyTorch实现重写/改写Dataset并载入Dataloader的完整攻略。

1. Dataset的重写/改写

1.1 创建自定义的Dataset

使用PyTorch构建Dataset需要继承torch.utils.data.Dataset类，并重新实现__init__、__len__、__getitem__三个方法。其中，__init__方法用于实现数据集初始化，__len__方法用于返回数据集的总长度，__getitem__方法用于通过索引获取数据。

from torch.utils.data import Dataset

class MyDataset(Dataset):
   def __init__(self, data_path):
       self.data_path = data_path
       # TODO: 初始化数据集

   def __len__(self):
       # TODO: 返回数据集的总长度
       return len(self.data)

   def __getitem__(self, index):
       # TODO: 通过索引获取数据
       return self.data[index]

1.2 自定义数据集的读取方式

默认情况下，PyTorch的Dataset读取数据的方式是使用PIL.Image.open，但是如果你的数据存储格式不同，你需要对读取方式进行修改。

from PIL import Image
from torch.utils.data import Dataset

class MyDataset(Dataset):
   def __init__(self, data_path):
       self.data_path = data_path
       # TODO: 初始化数据集

   def __len__(self):
       # TODO: 返回数据集的总长度
       return len(self.data)

   def __getitem__(self, index):
       # TODO: 通过索引获取数据
       img_path, label = self.data[index]
       img = Image.open(img_path).convert('RGB')
       return img, label

2. Dataloader的重写/改写

2.1 创建自定义的Dataloader

使用PyTorch构建Dataloader需要继承torch.utils.data.DataLoader类，并重新实现__init__方法。其中，__init__方法用于实现数据集初始化，包括数据集的载入方式、batch size、shuffle等。

from torch.utils.data import DataLoader

class MyDataLoader(DataLoader):
   def __init__(self, dataset, batch_size=1, shuffle=False, num_workers=0, collate_fn=None):
       super(MyDataLoader, self).__init__(dataset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers,
                                    collate_fn=collate_fn)
       # TODO: 自定义初始化

2.2 自定义collate_fn

collate_fn是一个可选参数，用来指定对batch数据的预处理方式。默认情况下，它会将每个数据按照Dataset返回的方式拼接成一个batch，但是如果你的数据不是相同形状的，你需要自定义collate_fn，将不同形状的数据拼接成相同形状的batch。

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

class MyDataset(Dataset):
    def __init__(self, data_path):
        self.data_path = data_path
        # TODO: 初始化数据集

    def __len__(self):
        # TODO: 返回数据集的总长度
        pass

    def __getitem__(self, index):
        # TODO: 通过索引获取数据
        pass

def collate_fn(batch):
    imgs = []
    labels = []
    for sample in batch:
        img, label = sample
        img = transforms.Resize((224, 224))(img)  # 将图像转换为指定大小
        img = transforms.ToTensor()(img)  # 将图像转换为Tensor
        imgs.append(img)
        labels.append(label)
    return torch.stack(imgs, 0), torch.tensor(labels)

dataset = MyDataset(data_path)
data_loader = DataLoader(dataset, batch_size=16, shuffle=True, num_workers=4, collate_fn=collate_fn)

以上是自定义Dataset和Dataloader的代码示例。根据实际需求，你可以对这些代码进行修改和扩展，以实现自己的目标。