VIII. Data Loading & Preprocessing (数据加载与预处理)#

1. `torch.utils.data.Dataset`#

Abstract base class for custom datasets. Must implement __len__ and __getitem__.

1
from torch.utils.data import Dataset
2

3
class MyDataset(Dataset):
4
    def __init__(self, data, labels):
5
        self.data = data
6
        self.labels = labels
7

8
    def __len__(self):
9
        return len(self.data)
10

11
    def __getitem__(self, idx):
12
        return self.data[idx], self.labels[idx]

Note: Put all preprocessing / augmentation inside __getitem__ for Lazy Loading (懒加载).

2. `torch.utils.data.DataLoader`#

Wraps a Dataset into an iterable batch loader with parallel reading (并行读取) and data shuffling (数据打乱).

1
from torch.utils.data import DataLoader
2
loader = DataLoader(
3
    dataset=train_ds, batch_size=32,
4
    shuffle=True, num_workers=4, pin_memory=True
5
)
6
for x, y in loader:
7
    ...

Note: On Windows, num_workers > 0 requires if __name__ == '__main__': guard.

3. `torchvision.transforms`#

Image preprocessing and data augmentation (数据增强) library. Chain multiple transforms with Compose.

1
from torchvision import transforms
2
transform = transforms.Compose([
3
    transforms.Resize(256),
4
    transforms.RandomCrop(224),
5
    transforms.RandomHorizontalFlip(),
6
    transforms.ToTensor(),
7
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
8
                         std=[0.229, 0.224, 0.225])
9
])

Note: The Normalize parameters are ImageNet statistics. Keep them consistent when using Transfer Learning (迁移学习).

4. `torchvision.datasets.ImageFolder`#

Automatically builds an image classification dataset from directory structure — subdirectory names become class labels (类别标签).

1
from torchvision.datasets import ImageFolder
2
# data/train/cat/*.jpg, data/train/dog/*.jpg
3
ds = ImageFolder(root='data/train', transform=transform)
4
print(ds.classes)   # ['cat', 'dog']

Note: Save the class_to_idx dictionary alongside the model checkpoint.

5. `torch.utils.data.random_split()`#

Randomly splits a dataset into train/validation subsets by specified lengths.

1
from torch.utils.data import random_split
2
n_val = int(len(dataset) * 0.2)
3
train_ds, val_ds = random_split(dataset, [len(dataset) - n_val, n_val])

Note: Pass generator=torch.Generator().manual_seed(42) for reproducible splits.

6. `torchvision.models` (pretrained)#

Provides many pre-trained models: ResNet, VGG, ViT, etc. Enables rapid Transfer Learning (迁移学习).

1
import torchvision.models as models
2
model = models.resnet50(weights='IMAGENET1K_V2')
3
model.fc = nn.Linear(2048, 10)  # replace head for fine-tuning

Note: Freeze the backbone: for p in model.parameters(): p.requires_grad = False.

💡 One-line Takeaway
The data pipeline is: Dataset (what) → transforms (how to augment) → DataLoader (how to batch).

VIII. Data Loading & Preprocessing (数据加载与预处理)#

1. torch.utils.data.Dataset#

2. torch.utils.data.DataLoader#

3. torchvision.transforms#

4. torchvision.datasets.ImageFolder#

5. torch.utils.data.random_split()#

6. torchvision.models (pretrained)#