utils.data模块

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

# 为了让示例可独立运行，创建假的数据集
def create_dummy_dataset(root_dir="dummy_data", num_samples=20):
    """创建一个假的图像数据集用于演示"""
    img_dir = os.path.join(root_dir, "images")
    if not os.path.exists(img_dir):
        os.makedirs(img_dir)

    # 创建假的CSV文件
    labels = []
    for i in range(num_samples):
        filename = f"{i:03d}.png"
        label = "cat" if i % 2 == 0 else "dog"
        labels.append({"image_filename": filename, "label": label})
        
        # 创建假的图像文件 (10x10的随机像素)
        dummy_img = Image.fromarray((torch.rand(10, 10, 3) * 255).numpy().astype('uint8'))
        dummy_img.save(os.path.join(img_dir, filename))

    df = pd.DataFrame(labels)
    df.to_csv(os.path.join(root_dir, "labels.csv"), index=False)
    print(f"Dummy dataset created at '{root_dir}'")

create_dummy_dataset()
# ==========================================================================================

# 1. 定义自定义 Dataset 类，继承 torch.utils.data.Dataset
class CustomImageDataset(Dataset):
    # (1) __init__: 初始化，加载数据清单
    def __init__(self, annotations_file, img_dir, transform=None, target_transform=None):
        """
        Args:
            annotations_file (string): 包含标签的CSV文件路径
            img_dir (string): 包含所有图像的目录路径
            transform (callable, optional): 应用于图像样本的可选转换
            target_transform (callable, optional): 应用于标签的可选转换
        """
        # 读取CSV文件
        self.img_labels = pd.read_csv(annotations_file)
        self.img_dir = img_dir
        self.transform = transform
        self.target_transform = target_transform
        
        # 为了演示，将文本标签映射为整数
        self.class_map = {"cat": 0, "dog": 1}

    # (2) __len__: 返回数据集的总大小
    def __len__(self):
        return len(self.img_labels)

    # (3) __getitem__: 根据索引加载并返回一个样本
    def __getitem__(self, idx):
        # 构造图像文件的完整路径
        img_path = os.path.join(self.img_dir, self.img_labels.iloc[idx, 0])
        
        # 加载图像
        image = Image.open(img_path).convert("RGB") # 确保图像是RGB格式
        
        # 获取标签
        label_text = self.img_labels.iloc[idx, 1]
        label = self.class_map[label_text] # 转换为整数标签

        # 应用转换
        if self.transform:
            image = self.transform(image)
        if self.target_transform:
            label = self.target_transform(label)
            
        return image, label

# 2. 实例化自定义 Dataset
data_root = "dummy_data"
csv_file = os.path.join(data_root, "labels.csv")
image_directory = os.path.join(data_root, "images")

# 定义一些图像转换
# 调整大小 -> 转换为Tensor -> 标准化
image_transforms = transforms.Compose([
    transforms.Resize((32, 32)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])

# 创建数据集实例
custom_dataset = CustomImageDataset(
    annotations_file=csv_file,
    img_dir=image_directory,
    transform=image_transforms
)

# 3. 验证 Dataset 的功能
print(f"\nDataset size: {len(custom_dataset)}") # 调用 __len__

# 获取第一个样本
first_image, first_label = custom_dataset[0] # 调用 __getitem__(0)
print(f"\nFirst sample - Image shape: {first_image.shape}, Label: {first_label}")
# 检查数据类型
print(f"Image dtype: {first_image.dtype}, Label type: {type(first_label)}")

# 4. 将 Dataset 传递给 DataLoader
# dataset的最终目的就是传给dataloader加载数据
data_loader = DataLoader(
    dataset=custom_dataset,
    batch_size=4,
    shuffle=True,
    num_workers=0 # 在Windows上，多进程需要在 main 函数里使用，先设为0
)

# 从 DataLoader 中获取一个批次的数据
# 这会在内部多次调用 custom_dataset.__getitem__()
images_batch, labels_batch = next(iter(data_loader))
print(f"\nBatch of images shape: {images_batch.shape}") # (batch_size, C, H, W)
print(f"Batch of labels shape: {labels_batch.shape}")
print(f"Labels in the batch: {labels_batch}")
  

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

# 重用上一节的 Dataset 定义和虚拟数据创建函数
def create_dummy_dataset(root_dir="dummy_data", num_samples=50): # 增加样本数以便演示
    img_dir = os.path.join(root_dir, "images")
    if not os.path.exists(img_dir): os.makedirs(img_dir)
    labels = [{"image_filename": f"{i:03d}.png", "label": "cat" if i % 2 == 0 else "dog"} for i in range(num_samples)]
    for item in labels:
        dummy_img = Image.fromarray((torch.rand(10, 10, 3) * 255).numpy().astype('uint8'))
        dummy_img.save(os.path.join(img_dir, item["image_filename"]))
    df = pd.DataFrame(labels)
    df.to_csv(os.path.join(root_dir, "labels.csv"), index=False)
    print(f"Dummy dataset created at '{root_dir}'")

class CustomImageDataset(Dataset):
    def __init__(self, annotations_file, img_dir, transform=None):
        self.img_labels = pd.read_csv(annotations_file)
        self.img_dir = img_dir
        self.transform = transform
        self.class_map = {"cat": 0, "dog": 1}
    def __len__(self):
        return len(self.img_labels)
    def __getitem__(self, idx):
        img_path = os.path.join(self.img_dir, self.img_labels.iloc[idx, 0])
        image = Image.open(img_path).convert("RGB")
        label = self.class_map[self.img_labels.iloc[idx, 1]]
        if self.transform:
            image = self.transform(image)
        return image, label

# ------------------------------------------------------------------------------------
# 主执行代码 - 必须放在 if __name__ == '__main__': 保护块中，以便安全使用 num_workers
# -------------------------------------------------------------------------------------
if __name__ == '__main__':
    # 1. 创建虚拟数据集
    create_dummy_dataset()

    # 2. 实例化 Dataset
    image_transforms = transforms.Compose([
        transforms.Resize((32, 32)),
        transforms.ToTensor()
    ])
    full_dataset = CustomImageDataset(
        annotations_file="dummy_data/labels.csv",
        img_dir="dummy_data/images",
        transform=image_transforms
    )

    # 3. 实例化 DataLoader (分别为训练和验证)
    # 训练 DataLoader: 打乱数据，使用多进程
    train_dataloader = DataLoader(
        dataset=full_dataset,
        batch_size=8,
        shuffle=True,      # 打乱数据
        num_workers=2,     # 使用2个子进程加载数据
        pin_memory=True,   # 如果使用GPU，建议开启
        drop_last=True     # 丢弃最后一个不完整的批次
    )

    # 验证 DataLoader: 不打乱，通常可以使用更大的 batch_size
    val_dataloader = DataLoader(
        dataset=full_dataset,
        batch_size=16,
        shuffle=False,     # 验证时不需要打乱
        num_workers=0      # 验证时数据量小，0也可以
    )

    # 4. 演示如何使用 DataLoader
    print(f"\nTotal samples in dataset: {len(full_dataset)}")
    
    # 训练 DataLoader 演示
    print("\n--- Iterating through train_dataloader ---")
    print(f"Number of batches in train_dataloader: {len(train_dataloader)}")
    
    # 从 DataLoader 中获取一个批次
    # 这是一个典型的训练循环的开始
    images_batch, labels_batch = next(iter(train_dataloader))
    
    print(f"Shape of one batch of images: {images_batch.shape}") # [batch_size, C, H, W]
    print(f"Shape of one batch of labels: {labels_batch.shape}")
    print(f"Labels in this batch: {labels_batch.numpy()}")
    
    # 模拟一个训练 epoch
    print("\nSimulating a training epoch loop...")
    for batch_idx, (data, targets) in enumerate(train_dataloader):
        # 在这里，你会将 data 和 targets 移动到 GPU
        # data, targets = data.to(device), targets.to(device)
        # 然后进行模型的前向传播、计算损失、反向传播...
        if batch_idx < 2: # 只打印前两个批次的信息
            print(f"  Batch {batch_idx+1}: Image batch shape {data.shape}, Label batch shape {targets.shape}")
    print("Training epoch simulation finished.")