Pytorch – データセットを学習用、テスト用に分割する方法

import torch
import torchvision.datasets as datasets

# 元のデータセット
dataset_dir = "/data/hymenoptera_data/train"
full_dataset = datasets.ImageFolder(dataset_dir)

# 学習データ、検証データに 8:2 の割合で分割する。
train_size = int(0.8 * len(full_dataset))
val_size = len(full_dataset) - train_size
train_dataset, val_dataset = torch.utils.data.random_split(
    full_dataset, [train_size, val_size]
)

print(f"full: {len(full_dataset)} -> train: {len(train_dataset)}, test: {len(val_dataset)}")

import numpy as np
import torch
import torchvision.datasets as datasets

# 元のデータセット
dataset_dir = "/data/hymenoptera_data/train"
dataset = datasets.ImageFolder(dataset_dir)

# 学習データ、検証データに 8:2 の割合で分割する。
train_size = int(0.8 * len(full_dataset))
indices = np.arange(len(full_dataset))

train_dataset = torch.utils.data.Subset(dataset, indices[:train_size])
val_dataset = torch.utils.data.Subset(dataset, indices[train_size:])

print(f"full: {len(dataset)} -> train: {len(train_dataset)}, test: {len(val_dataset)}")

import numpy as np
import torch
import torchvision.datasets as datasets
import torchvision.transforms as transforms

train_transform = transforms.Compose(
    [
        transforms.RandomResizedCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ]
)

val_transform = transforms.Compose(
    [
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ]
)

class MySubset(torch.utils.data.Dataset):
    def __init__(self, dataset, indices, transform=None):
        self.dataset = dataset
        self.indices = indices
        self.transform = transform

    def __getitem__(self, idx):
        img, label = self.dataset[self.indices[idx]]
        if self.transform:
            img = self.transform(img)

        return img, label

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


# 元のデータセット
dataset_dir = "/data/hymenoptera_data/train"
dataset = datasets.ImageFolder(dataset_dir)

# 学習データ、検証データに 8:2 の割合で分割する。
train_size = int(0.8 * len(full_dataset))
indices = np.arange(len(full_dataset))

train_dataset = MySubset(dataset, indices[:train_size], train_transform)
val_dataset = MySubset(dataset, indices[train_size:], val_transform)

print(f"full: {len(dataset)} -> train: {len(train_dataset)}, test: {len(val_dataset)}")