05-PyTorch自定义数据集Datasets、Loader和tranform

import torch
from torch import nn

device = "cuda" if torch.cuda.is_available() else "cpu"

import requests
import zipfile
from pathlib import Path

# Setup path to data folder
data_path = Path("data/")
image_path = data_path / "pizza_steak_sushi"

# If the image folder doesn't exist, download it and prepare it... 
if image_path.is_dir():
    print(f"{image_path} directory exists.")
else:
    print(f"Did not find {image_path} directory, creating one...")
    image_path.mkdir(parents=True, exist_ok=True)
    
    # Download pizza, steak, sushi data
    with open(data_path / "pizza_steak_sushi.zip", "wb") as f:
        request = requests.get("https://github.com/mrdbourke/pytorch-deep-learning/raw/main/data/pizza_steak_sushi.zip")
        print("Downloading pizza, steak, sushi data...")
        f.write(request.content)

    # Unzip pizza, steak, sushi data
    with zipfile.ZipFile(data_path / "pizza_steak_sushi.zip", "r") as zip_ref:
        print("Unzipping pizza, steak, sushi data...") 
        zip_ref.extractall(image_path)

pizza_steak_sushi/ <- overall dataset folder
    train/ <- training images
        pizza/ <- class name as folder name
            image01.jpeg
            image02.jpeg
            ...
        steak/
            image24.jpeg
            image25.jpeg
            ...
        sushi/
            image37.jpeg
            ...
    test/ <- testing images
        pizza/
            image101.jpeg
            image102.jpeg
            ...
        steak/
            image154.jpeg
            image155.jpeg
            ...
        sushi/
            image167.jpeg
            ...

# Setup train and testing paths
train_dir = image_path / "train"
test_dir = image_path / "test"

train_dir, test_dir

(PosixPath('data/pizza_steak_sushi/train'),
 PosixPath('data/pizza_steak_sushi/test'))

import random
from PIL import Image

# 设置随机种子
random.seed(42)  

# 1. 获取所有图像路径（*表示"任意组合"）
image_path_list = list(image_path.glob("*/*/*.jpg")) # 用 pathlib.Path.glob() 获取所有图像路径，以查找以 .jpg 结尾的所有文件。

# 2. 获取随机图像路径
random_image_path = random.choice(image_path_list)

# 3. 从路径名中获取图像类别（图像类别是存储图像的目录名称）
image_class = random_image_path.parent.stem 

# 4. 打开图像
img = Image.open(random_image_path)

# 5. 打印元数据
print(f"随机图像路径：{random_image_path}")
print(f"图像类别：{image_class}")
print(f"图像高度：{img.height}")
print(f"图像宽度：{img.width}")
img

Random image path: data/pizza_steak_sushi/test/pizza/2124579.jpg
Image class: pizza
Image height: 384
Image width: 512

import numpy as np
import matplotlib.pyplot as plt

# Turn the image into an array
img_as_array = np.asarray(img)

# Plot the image with matplotlib
plt.figure(figsize=(10, 7))
plt.imshow(img_as_array)
plt.title(f"Image class: {image_class} | Image shape: {img_as_array.shape} -> [height, width, color_channels]")
plt.axis(False);

import torch
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

# Write transform for image
data_transform = transforms.Compose([
    # 整图像大小（从大约 512x512 到 64x64)
    transforms.Resize(size=(64, 64)),
    # 水平方向上随机翻转图像
    transforms.RandomHorizontalFlip(p=0.5), # p = probability of flip, 0.5 = 50% chance
    # 从 PIL 图像转换为 PyTorch 张量。
    transforms.ToTensor() #  0 到 255 -> 0.0 到 1.0 
])

def plot_transformed_images(image_paths, transform, n=3, seed=42):
    """从图像路径中绘制一系列随机图像。

        将从image_paths中打开n个图像路径，使用transform进行转换，并将它们并排绘制出来。

        参数：
        image_paths（列表）：目标图像路径列表。
        transform（PyTorch转换）：要应用于图像的转换。
        n（整数，可选）：要绘制的图像数量。默认为3。
        seed（整数，可选）：随机生成器的随机种子。默认为42。
    """
    random.seed(seed)
    random_image_paths = random.sample(image_paths, k=n)
    for image_path in random_image_paths:
        with Image.open(image_path) as f:
            fig, ax = plt.subplots(1, 2)
            ax[0].imshow(f) 
            ax[0].set_title(f"Original \nSize: {f.size}")
            ax[0].axis("off")

            #转换函数并且绘制图形
            # Note: permute() 重新图像排列顺序：
            # (PyTorch default is [C, H, W] but Matplotlib is [H, W, C])
            transformed_image = transform(f).permute(1, 2, 0) 
            ax[1].imshow(transformed_image) 
            ax[1].set_title(f"Transformed \nSize: {transformed_image.shape}")
            ax[1].axis("off")

            fig.suptitle(f"Class: {image_path.parent.stem}", fontsize=16)

plot_transformed_images(image_path_list, 
                        transform=data_transform, 
                        n=3)

# 使用ImageFolder创建数据集
from torchvision import datasets
train_data = datasets.ImageFolder(root=train_dir,  # 图像的目标文件夹
                                  transform=data_transform,  # 对数据（图像）执行的转换
                                  target_transform=None)  # 对标签执行的转换（如果需要的话）

test_data = datasets.ImageFolder(root=test_dir,
                                 transform=data_transform)

print(f"1.训练数据：\n{train_data}\n2.测试数据：\n{test_data}")

1.训练数据：
Dataset ImageFolder
    Number of datapoints: 225
    Root location: data/pizza_steak_sushi/train
    StandardTransform
Transform: Compose(
               Resize(size=(64, 64), interpolation=bilinear, max_size=None, antialias=None)
               RandomHorizontalFlip(p=0.5)
               ToTensor()
           )
2.测试数据：
Dataset ImageFolder
    Number of datapoints: 75
    Root location: data/pizza_steak_sushi/test
    StandardTransform
Transform: Compose(
               Resize(size=(64, 64), interpolation=bilinear, max_size=None, antialias=None)
               RandomHorizontalFlip(p=0.5)
               ToTensor()
           )

#数据集的类别
class_names = train_data.classes
class_names
>>>
['pizza', 'steak', 'sushi']

# 数据集的类别的字典形式
class_dict = train_data.class_to_idx
class_dict
>>>
{'pizza': 0, 'steak': 1, 'sushi': 2}

# 数据集长度
len(train_data), len(test_data)
>>>
(225, 75)

img, label = train_data[0][0], train_data[0][1]
print(f"Image tensor:\n{img}")
print(f"Image shape: {img.shape}")
print(f"Image datatype: {img.dtype}")
print(f"Image label: {label}")
print(f"Label datatype: {type(label)}")
>>>
Image tensor:
tensor([[[0.1137, 0.1020, 0.0980,  ..., 0.1255, 0.1216, 0.1176],
         [0.1059, 0.0980, 0.0980,  ..., 0.1294, 0.1294, 0.1294],
         [0.1020, 0.0980, 0.0941,  ..., 0.1333, 0.1333, 0.1333],
         ...,
         [0.1098, 0.1098, 0.1255,  ..., 0.1686, 0.1647, 0.1686],
         [0.0863, 0.0941, 0.1098,  ..., 0.1686, 0.1647, 0.1686],
         [0.0863, 0.0863, 0.0980,  ..., 0.1686, 0.1647, 0.1647]],

        [[0.0745, 0.0706, 0.0745,  ..., 0.0588, 0.0588, 0.0588],
         [0.0706, 0.0706, 0.0745,  ..., 0.0627, 0.0627, 0.0627],
         [0.0706, 0.0745, 0.0745,  ..., 0.0706, 0.0706, 0.0706],
         ...,
         [0.1255, 0.1333, 0.1373,  ..., 0.2510, 0.2392, 0.2392],
         [0.1098, 0.1176, 0.1255,  ..., 0.2510, 0.2392, 0.2314],
         [0.1020, 0.1059, 0.1137,  ..., 0.2431, 0.2353, 0.2275]],

        [[0.0941, 0.0902, 0.0902,  ..., 0.0196, 0.0196, 0.0196],
         [0.0902, 0.0863, 0.0902,  ..., 0.0196, 0.0157, 0.0196],
         [0.0902, 0.0902, 0.0902,  ..., 0.0157, 0.0157, 0.0196],
         ...,
         [0.1294, 0.1333, 0.1490,  ..., 0.1961, 0.1882, 0.1804],
         [0.1098, 0.1137, 0.1255,  ..., 0.1922, 0.1843, 0.1804],
         [0.1059, 0.1020, 0.1059,  ..., 0.1843, 0.1804, 0.1765]]])
Image shape: torch.Size([3, 64, 64])
Image datatype: torch.float32
Image label: 0
Label datatype: <class 'int'>

# 重新排列图像维度
img_permute = img.permute(1, 2, 0)

# 打印差别
print(f"Original shape: {img.shape} -> [color_channels, height, width]")
print(f"Image permute shape: {img_permute.shape} -> [height, width, color_channels]")

# 绘制
plt.figure(figsize=(10, 7))
plt.imshow(img.permute(1, 2, 0))
plt.axis("off")
plt.title(class_names[label], fontsize=14);

Original shape: torch.Size([3, 64, 64]) -> [color_channels, height, width]
Image permute shape: torch.Size([64, 64, 3]) -> [height, width, color_channels]

# 将加载的Dataset 转为 DataLoader
from torch.utils.data import DataLoader
train_dataloader = DataLoader(dataset=train_data, 
                              batch_size=1,  # 每批次多少个数据
                              num_workers=1, #  ow many subprocesses to use for data loading. 0 means that the data will be loaded in the main process. (default: 0)
                              shuffle=True) # 打乱数据

test_dataloader = DataLoader(dataset=test_data, 
                             batch_size=1, 
                             num_workers=1, 
                             shuffle=False) # 测试数据不需要打乱

train_dataloader, test_dataloader

(<torch.utils.data.dataloader.DataLoader at 0x7f53c0b9dca0>,
 <torch.utils.data.dataloader.DataLoader at 0x7f53c0b9de50>)

img, label = next(iter(train_dataloader))

# Batch size will now be 1, try changing the batch_size parameter above and see what happens
print(f"Image shape: {img.shape} -> [batch_size, color_channels, height, width]")
print(f"Label shape: {label.shape}")

Image shape: torch.Size([1, 3, 64, 64]) -> [batch_size, color_channels, height, width]
Label shape: torch.Size([1])

import os
import pathlib
import torch

from PIL import Image
from torch.utils.data import Dataset
from torchvision import transforms
from typing import Tuple, Dict, List

# 创建函数以在目标目录中查找类别
def find_classes(directory: str) -> Tuple[List[str], Dict[str, int]]:
    """在目标目录中查找类别文件夹名称。
    
    假设目标目录采用标准的图像分类格式。

    参数:
        directory (str): 要从中加载类别名称的目标目录。

    返回:
        Tuple[List[str], Dict[str, int]]: (类别名称列表, 类别名称: 索引...的字典)
    
    示例:
        find_classes("food_images/train")
        >>> (["class_1", "class_2"], {"class_1": 0, ...})
    """
    # 1. 通过扫描目标目录获取类别名称
    classes = sorted(entry.name for entry in os.scandir(directory) if entry.is_dir())
    
    # 2. 如果找不到类别名称，则引发错误
    if not classes:
        raise FileNotFoundError(f"在{directory}中找不到任何类别。")
        
    # 3. 创建索引标签的字典（计算机更喜欢数字标签而不是字符串标签）
    class_to_idx = {cls_name: i for i, cls_name in enumerate(classes)}
    return classes, class_to_idx

find_classes(train_dir)

(['pizza', 'steak', 'sushi'], {'pizza': 0, 'steak': 1, 'sushi': 2})

# 编写自定义数据集类（继承自torch.utils.data.Dataset）
from torch.utils.data import Dataset

# 1. 子类化torch.utils.data.Dataset
class ImageFolderCustom(Dataset):
    
    # 2. 使用targ_dir和transform（可选）参数进行初始化
    def __init__(self, targ_dir: str, transform=None) -> None: # transform可以不填，后文会定义一个transform
        
        # 3. 定义类属性 此处根据需求可以进行修改
        # a.获取所有图像路径
        self.paths = list(pathlib.Path(targ_dir).glob("*/*.jpg"))  # 注意：如果有.png或.jpeg文件，您需要更新此处
        # b.设置转换
        self.transform = transform
        # c.创建classes和class_to_idx属性
        self.classes, self.class_to_idx = find_classes(targ_dir) # 在5.1中定义过 返回类和对应的字典

    # 4. 创建加载图像的函数
    def load_image(self, index: int) -> Image.Image:
        "通过路径打开图像并返回它。"
        image_path = self.paths[index]
        return Image.open(image_path) 
    
    # 5. 覆写__len__()方法（可选，但建议创建）
    def __len__(self) -> int:
        "返回样本的总数。"
        return len(self.paths) # 直接调用len函数计算self.paths的长度
    
    # 6. 覆写__getitem__()方法（torch.utils.data.Dataset的子类所必需）
    def __getitem__(self, index: int) -> Tuple[torch.Tensor, int]:
        "返回一个数据样本，数据和标签（X, y）。"
        img = self.load_image(index)
        class_name  = self.paths[index].parent.name  # 期望路径为data_folder/class_name/image.jpeg
        class_idx = self.class_to_idx[class_name]

        # 如果需要，进行转换
        if self.transform:
            return self.transform(img), class_idx  # 返回数据和标签（X, y）
        else:
            return img, class_idx  # 返回数据和标签（X, y）

# 1 创建transform函数
# 转换训练数据
train_transforms = transforms.Compose([
    transforms.Resize((64, 64)),
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.ToTensor()
])

# 转换测试数据
test_transforms = transforms.Compose([
    transforms.Resize((64, 64)),
    transforms.ToTensor()
])

# 2. 实例化训练和测试数据集 
train_data_custom = ImageFolderCustom(targ_dir=train_dir, 
                                      transform=train_transforms)
test_data_custom = ImageFolderCustom(targ_dir=test_dir, 
                                     transform=test_transforms)
train_data_custom, test_data_custom

(<__main__.ImageFolderCustom at 0x7f5461f70c70>,
 <__main__.ImageFolderCustom at 0x7f5461f70c40>)

len(train_data_custom), len(test_data_custom)
>>> (225, 75)

train_data_custom.classes
>>> ['pizza', 'steak', 'sushi']

train_data_custom.class_to_idx
>>> {'pizza': 0, 'steak': 1, 'sushi': 2}

from torch.utils.data import DataLoader
train_dataloader_custom = DataLoader(dataset=train_data_custom,
                                     batch_size=1, 
                                     num_workers=0, 
                                     shuffle=True) 

test_dataloader_custom = DataLoader(dataset=test_data_custom, 
                                    batch_size=1, 
                                    num_workers=0, 
                                    shuffle=False) 

from torchvision import transforms

train_transforms = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.TrivialAugmentWide(num_magnitude_bins=31), # how intense 
    transforms.ToTensor() 
])

# 通常不会在测试集上执行数据增强。数据增强的想法是人为地增加训练集的多样性，以更好地对测试集进行预测。
test_transforms = transforms.Compose([
    transforms.Resize((224, 224)), 
    transforms.ToTensor()
])

image_path_list = list(image_path.glob("*/*/*.jpg"))

plot_transformed_images(
    image_paths=image_path_list,
    transform=train_transforms,
    n=3,
    seed=None

# Create simple transform
simple_transform = transforms.Compose([ 
    transforms.Resize((64, 64)),
    transforms.ToTensor(),
])

# 1. 加载data
from torchvision import datasets
train_data_simple = datasets.ImageFolder(root=train_dir, transform=simple_transform)
test_data_simple = datasets.ImageFolder(root=test_dir, transform=simple_transform)

# 2. 将数据转化为DataLoaders
import os
from torch.utils.data import DataLoader


BATCH_SIZE = 32
NUM_WORKERS = os.cpu_count()
print(f"使用 {BATCH_SIZE}批次大小 和 {NUM_WORKERS} 个工作线程创建DataLoaders")

#  创建DataLoaders
train_dataloader_simple = DataLoader(train_data_simple, 
                                     batch_size=BATCH_SIZE, 
                                     shuffle=True, 
                                     num_workers=NUM_WORKERS)

test_dataloader_simple = DataLoader(test_data_simple, 
                                    batch_size=BATCH_SIZE, 
                                    shuffle=False, 
                                    num_workers=NUM_WORKERS)

train_dataloader_simple, test_dataloader_simple

class TinyVGG(nn.Module):
    """
    模型架构，参考自：https://poloclub.github.io/cnn-explainer/
    """
    def __init__(self, input_shape: int, hidden_units: int, output_shape: int) -> None:
        super().__init__()
        self.conv_block_1 = nn.Sequential(
            nn.Conv2d(in_channels=input_shape, 
                      out_channels=hidden_units, 
                      kernel_size=3,  # 内核的的大小是3*3
                      stride=1,  # 步长，默认值
                      padding=1),  # 选项 = "valid"（无填充）或 "same"（输出与输入形状相同）或整数表示特定的数值
            nn.ReLU(),
            nn.Conv2d(in_channels=hidden_units, 
                      out_channels=hidden_units,
                      kernel_size=3,
                      stride=1,
                      padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2,
                         stride=2)  # 默认的 stride 值与 kernel_size 相同
        )
        self.conv_block_2 = nn.Sequential(
            nn.Conv2d(hidden_units, hidden_units, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv2d(hidden_units, hidden_units, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            # 这个 in_features 形状是从哪里来的？
            # 这是因为网络的每一层都会压缩和改变输入数据的形状。
            nn.Linear(in_features=hidden_units*16*16,
                      out_features=output_shape)
        )
    
    def forward(self, x: torch.Tensor):
        x = self.conv_block_1(x)
        # print(x.shape)
        x = self.conv_block_2(x)
        # print(x.shape)
        x = self.classifier(x)
        # print(x.shape)
        return x
        # return self.classifier(self.conv_block_2(self.conv_block_1(x)))  # <- 利用运算符融合的好处

torch.manual_seed(42)
model_0 = TinyVGG(input_shape=3,  # 颜色通道数（RGB为3）
                  hidden_units=10, 
                  output_shape=len(train_data.classes)).to(device)
model_0

TinyVGG(
  (conv_block_1): Sequential(
    (0): Conv2d(3, 10, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU()
    (2): Conv2d(10, 10, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (3): ReLU()
    (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (conv_block_2): Sequential(
    (0): Conv2d(10, 10, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU()
    (2): Conv2d(10, 10, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (3): ReLU()
    (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (classifier): Sequential(
    (0): Flatten(start_dim=1, end_dim=-1)
    (1): Linear(in_features=2560, out_features=3, bias=True)
  )
)

pip install torchinfo

from torchinfo import summary
summary(model_0, input_size=[1, 3, 64, 64]) 

def train_step(model: torch.nn.Module, 
               dataloader: torch.utils.data.DataLoader, 
               loss_fn: torch.nn.Module, 
               optimizer: torch.optim.Optimizer):
    # 将模型设为训练模式
    model.train()
    
    # 设置训练损失和训练准确率的初始值
    train_loss, train_acc = 0, 0
    
    # 遍历数据加载器中的数据批次
    for batch, (X, y) in enumerate(dataloader):
        # 将数据发送到目标设备
        X, y = X.to(device), y.to(device)

        # 1. 正向传播
        y_pred = model(X)

        # 2. 计算并累积损失
        loss = loss_fn(y_pred, y)
        train_loss += loss.item() 

        # 3. 优化器梯度清零
        optimizer.zero_grad()

        # 4. 损失反向传播
        loss.backward()

        # 5. 优化器更新参数
        optimizer.step()

        # 计算并累积准确率指标（在所有批次上）
        y_pred_class = torch.argmax(torch.softmax(y_pred, dim=1), dim=1)
        train_acc += (y_pred_class == y).sum().item()/len(y_pred)

    # 调整指标以得到每个批次的平均损失和准确率
    train_loss = train_loss / len(dataloader)
    train_acc = train_acc / len(dataloader)
    return train_loss, train_acc

def test_step(model: torch.nn.Module, 
              dataloader: torch.utils.data.DataLoader, 
              loss_fn: torch.nn.Module):
    # 将模型设为评估模式
    model.eval() 
    
    # 设置测试损失和测试准确率的初始值
    test_loss, test_acc = 0, 0
    
    # 打开推理上下文管理器
    with torch.inference_mode():
        # 遍历数据加载器中的数据批次
        for batch, (X, y) in enumerate(dataloader):
            # 将数据发送到目标设备
            X, y = X.to(device), y.to(device)
    
            # 1. 正向传播
            test_pred_logits = model(X)

            # 2. 计算并累积损失
            loss = loss_fn(test_pred_logits, y)
            test_loss += loss.item()
            
            # 计算并累积准确率
            test_pred_labels = test_pred_logits.argmax(dim=1)
            test_acc += ((test_pred_labels == y).sum().item()/len(test_pred_labels))
            
    # 调整指标以得到每个批次的平均损失和准确率
    test_loss = test_loss / len(dataloader)
    test_acc = test_acc / len(dataloader)
    return test_loss, test_acc

from tqdm.auto import tqdm

def train(model: torch.nn.Module, 
          train_dataloader: torch.utils.data.DataLoader, 
          test_dataloader: torch.utils.data.DataLoader, 
          optimizer: torch.optim.Optimizer,
          loss_fn: torch.nn.Module = nn.CrossEntropyLoss(),
          epochs: int = 5):

    # 2. 创造空字典储存结果
    results = {"train_loss": [],
        "train_acc": [],
        "test_loss": [],
        "test_acc": []
    }

    # 3. 循环训练和测试
    for epoch in tqdm(range(epochs)):
        train_loss, train_acc = train_step(model=model,
                                           dataloader=train_dataloader,
                                           loss_fn=loss_fn,
                                           optimizer=optimizer)
        test_loss, test_acc = test_step(model=model,
            dataloader=test_dataloader,
            loss_fn=loss_fn)

        # 4.打印训练成果
        print(
            f"Epoch: {epoch+1} | "
            f"train_loss: {train_loss:.4f} | "
            f"train_acc: {train_acc:.4f} | "
            f"test_loss: {test_loss:.4f} | "
            f"test_acc: {test_acc:.4f}"
        )

        # 5. 更新结果字典
        results["train_loss"].append(train_loss)
        results["train_acc"].append(train_acc)
        results["test_loss"].append(test_loss)
        results["test_acc"].append(test_acc)

    # 6. 在epoch结束时返回填充的结果
    return results

# 设置随机种子
torch.manual_seed(42) 
torch.cuda.manual_seed(42)

# 设置迭代次数
NUM_EPOCHS = 5

# 重新创建一个 TinyVGG 实例
model_0 = TinyVGG(input_shape=3,  # 颜色通道数（RGB 为 3）
                  hidden_units=10, 
                  output_shape=len(train_data.classes)).to(device)

# 设置损失函数和优化器
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(params=model_0.parameters(), lr=0.001)

# 启动计时器
from timeit import default_timer as timer 
start_time = timer()

# 训练 model_0
model_0_results = train(model=model_0, 
                        train_dataloader=train_dataloader_simple,
                        test_dataloader=test_dataloader_simple,
                        optimizer=optimizer,
                        loss_fn=loss_fn, 
                        epochs=NUM_EPOCHS)

# 停止计时器并打印所用时间
end_time = timer()
print(f"总训练时间：{end_time-start_time:.3f} 秒")

  0%|          | 0/5 [00:00<?, ?it/s]
Epoch: 1 | train_loss: 1.1078 | train_acc: 0.2578 | test_loss: 1.1360 | test_acc: 0.2604
Epoch: 2 | train_loss: 1.0847 | train_acc: 0.4258 | test_loss: 1.1620 | test_acc: 0.1979
Epoch: 3 | train_loss: 1.1157 | train_acc: 0.2930 | test_loss: 1.1697 | test_acc: 0.1979
Epoch: 4 | train_loss: 1.0956 | train_acc: 0.4141 | test_loss: 1.1384 | test_acc: 0.1979
Epoch: 5 | train_loss: 1.0985 | train_acc: 0.2930 | test_loss: 1.1426 | test_acc: 0.1979
Total training time: 4.935 seconds

# 检查 model_0_results 字典
model_0_results.keys()
>>> dict_keys(['train_loss', 'train_acc', 'test_loss', 'test_acc'])

def plot_loss_curves(results: Dict[str, List[float]]):
    """绘制结果字典中的训练曲线。

    参数:
        results (dict): 包含值列表的字典，例如：
            {"train_loss": [...],
             "train_acc": [...],
             "test_loss": [...],
             "test_acc": [...]}
    """
    
    # 获取结果字典中的损失值（训练和测试）
    loss = results['train_loss']
    test_loss = results['test_loss']

    # 获取结果字典中的准确率值（训练和测试）
    accuracy = results['train_acc']
    test_accuracy = results['test_acc']

    # 确定有多少个 epoch
    epochs = range(len(results['train_loss']))

    # 设置绘图
    plt.figure(figsize=(15, 7))

    # 绘制损失曲线
    plt.subplot(1, 2, 1)
    plt.plot(epochs, loss, label='train_loss')
    plt.plot(epochs, test_loss, label='test_loss')
    plt.title('损失')
    plt.xlabel('Epochs')
    plt.legend()

    # 绘制准确率曲线
    plt.subplot(1, 2, 2)
    plt.plot(epochs, accuracy, label='train_accuracy')
    plt.plot(epochs, test_accuracy, label='test_accuracy')
    plt.title('准确率')
    plt.xlabel('Epochs')
    plt.legend()

# 使用TrivialAugment创建训练集
train_transform_trivial_augment = transforms.Compose([
    transforms.Resize((64, 64)),
    transforms.TrivialAugmentWide(num_magnitude_bins=31),
    transforms.ToTensor() 
])

# 创建测试机
test_transform = transforms.Compose([
    transforms.Resize((64, 64)),
    transforms.ToTensor()
])

train_data_augmented = datasets.ImageFolder(train_dir, transform=train_transform_trivial_augment)
test_data_simple = datasets.ImageFolder(test_dir, transform=test_transform)

(Dataset ImageFolder
     Number of datapoints: 225
     Root location: data/pizza_steak_sushi/train
     StandardTransform
 Transform: Compose(
                Resize(size=(64, 64), interpolation=bilinear, max_size=None, antialias=None)
                TrivialAugmentWide(num_magnitude_bins=31, interpolation=InterpolationMode.NEAREST, fill=None)
                ToTensor()
            ),
 Dataset ImageFolder
     Number of datapoints: 75
     Root location: data/pizza_steak_sushi/test
     StandardTransform
 Transform: Compose(
                Resize(size=(64, 64), interpolation=bilinear, max_size=None, antialias=None)
                ToTensor()
            ))

# Turn Datasets into DataLoader's
import os
BATCH_SIZE = 32
NUM_WORKERS = os.cpu_count()

torch.manual_seed(42)
train_dataloader_augmented = DataLoader(train_data_augmented, 
                                        batch_size=BATCH_SIZE, 
                                        shuffle=True,
                                        num_workers=NUM_WORKERS)

test_dataloader_simple = DataLoader(test_data_simple, 
                                    batch_size=BATCH_SIZE, 
                                    shuffle=False, 
                                    num_workers=NUM_WORKERS)

# Create model_1 and send it to the target device
torch.manual_seed(42)
model_1 = TinyVGG(
    input_shape=3,
    hidden_units=10,
    output_shape=len(train_data_augmented.classes)).to(device)
model_1

TinyVGG(
  (conv_block_1): Sequential(
    (0): Conv2d(3, 10, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU()
    (2): Conv2d(10, 10, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (3): ReLU()
    (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (conv_block_2): Sequential(
    (0): Conv2d(10, 10, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU()
    (2): Conv2d(10, 10, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (3): ReLU()
    (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (classifier): Sequential(
    (0): Flatten(start_dim=1, end_dim=-1)
    (1): Linear(in_features=2560, out_features=3, bias=True)
  )
)

# Set random seeds
torch.manual_seed(42) 
torch.cuda.manual_seed(42)

# Set number of epochs
NUM_EPOCHS = 5

# Setup loss function and optimizer
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(params=model_1.parameters(), lr=0.001)

# Start the timer
from timeit import default_timer as timer 
start_time = timer()

# Train model_1
model_1_results = train(model=model_1, 
                        train_dataloader=train_dataloader_augmented,
                        test_dataloader=test_dataloader_simple,
                        optimizer=optimizer,
                        loss_fn=loss_fn, 
                        epochs=NUM_EPOCHS)

# End the timer and print out how long it took
end_time = timer()
print(f"Total training time: {end_time-start_time:.3f} seconds")

  0%|          | 0/5 [00:00<?, ?it/s]
Epoch: 1 | train_loss: 1.1074 | train_acc: 0.2500 | test_loss: 1.1058 | test_acc: 0.2604
Epoch: 2 | train_loss: 1.0791 | train_acc: 0.4258 | test_loss: 1.1382 | test_acc: 0.2604
Epoch: 3 | train_loss: 1.0803 | train_acc: 0.4258 | test_loss: 1.1685 | test_acc: 0.2604
Epoch: 4 | train_loss: 1.1285 | train_acc: 0.3047 | test_loss: 1.1623 | test_acc: 0.2604
Epoch: 5 | train_loss: 1.0880 | train_acc: 0.4258 | test_loss: 1.1472 | test_acc: 0.2604
Total training time: 4.924 seconds

plot_loss_curves(model_1_results)

import pandas as pd
model_0_df = pd.DataFrame(model_0_results)
model_1_df = pd.DataFrame(model_1_results)
model_0_df

# Setup a plot 
plt.figure(figsize=(15, 10))

# Get number of epochs
epochs = range(len(model_0_df))

# Plot train loss
plt.subplot(2, 2, 1)
plt.plot(epochs, model_0_df["train_loss"], label="Model 0")
plt.plot(epochs, model_1_df["train_loss"], label="Model 1")
plt.title("Train Loss")
plt.xlabel("Epochs")
plt.legend()

# Plot test loss
plt.subplot(2, 2, 2)
plt.plot(epochs, model_0_df["test_loss"], label="Model 0")
plt.plot(epochs, model_1_df["test_loss"], label="Model 1")
plt.title("Test Loss")
plt.xlabel("Epochs")
plt.legend()

# Plot train accuracy
plt.subplot(2, 2, 3)
plt.plot(epochs, model_0_df["train_acc"], label="Model 0")
plt.plot(epochs, model_1_df["train_acc"], label="Model 1")
plt.title("Train Accuracy")
plt.xlabel("Epochs")
plt.legend()

# Plot test accuracy
plt.subplot(2, 2, 4)
plt.plot(epochs, model_0_df["test_acc"], label="Model 0")
plt.plot(epochs, model_1_df["test_acc"], label="Model 1")
plt.title("Test Accuracy")
plt.xlabel("Epochs")
plt.legend();

# 下载自定义图像
import requests

# 设置自定义图像路径
custom_image_path = data_path / "04-pizza-dad.jpeg"


# 如果图像不存在，则下载图像
if not custom_image_path.is_file():
    with open(custom_image_path, "wb") as f:
        # 从 GitHub 下载时，需要使用 "raw" 文件链接
        request = requests.get("https://raw.githubusercontent.com/mrdbourke/pytorch-deep-learning/main/images/04-pizza-dad.jpeg")
        print(f"正在下载 {custom_image_path}...")
        f.write(request.content)
else:
    print(f"{custom_image_path} 已存在，跳过下载。")

# 加载自定义图像并将张量值转换为 float32 类型
custom_image = torchvision.io.read_image(str(custom_image_path)).type(torch.float32)

# 将图像像素值除以 255，使其范围在 [0, 1] 之间
custom_image = custom_image / 255. 

# 创建变换管道以调整图像大小
custom_image_transform = transforms.Compose([
    transforms.Resize((64, 64)),
])

# 对目标图像进行变换
custom_image_transformed = custom_image_transform(custom_image)

# 打印原始形状和新形状
# print(f"原始形状: {custom_image.shape}")
# print(f"新形状: {custom_image_transformed.shape}")

model_1.eval()
with torch.inference_mode():
    # 在图像上添加一个额外的维度
    custom_image_transformed_with_batch_size = custom_image_transformed.unsqueeze(dim=0)
    
    # 打印不同的形状
    print(f"自定义图像变换后的形状: {custom_image_transformed.shape}")
    print(f"添加维度后的自定义图像形状: {custom_image_transformed_with_batch_size.shape}")
    
    # 对带有额外维度的图像进行预测
    custom_image_pred = model_1(custom_image_transformed.unsqueeze(dim=0).to(device))

# 预测结果
custom_image_pred
>>>
tensor([[ 0.1172,  0.0160, -0.1425]], device='cuda:0')

# 打印预测的 logits
print(f"预测的 logits: {custom_image_pred}")

# 将 logits 转换为预测概率（使用 torch.softmax() 进行多类别分类）
custom_image_pred_probs = torch.softmax(custom_image_pred, dim=1)
print(f"预测的概率: {custom_image_pred_probs}")

# 将预测概率转换为预测标签
custom_image_pred_label = torch.argmax(custom_image_pred_probs, dim=1)
print(f"预测的标签: {custom_image_pred_label}")

预测的 logits:tensor([[ 0.1172,  0.0160, -0.1425]], device='cuda:0')
预测的概率: tensor([[0.3738, 0.3378, 0.2883]], device='cuda:0')
预测的标签: tensor([0], device='cuda:0')

# 找到类名
custom_image_pred_class = class_names[custom_image_pred_label.cpu()] 
custom_image_pred_class

def pred_and_plot_image(model: torch.nn.Module, 
                        image_path: str, 
                        class_names: List[str] = None, 
                        transform=None,
                        device: torch.device = device):
    """对目标图像进行预测并绘制图像及其预测结果。"""
    
    # 1. 加载图像并将张量值转换为 float32 类型
    target_image = torchvision.io.read_image(str(image_path)).type(torch.float32)
    
    # 2. 将图像的像素值除以 255，使其范围在 [0, 1] 之间
    target_image = target_image / 255. 
    
    # 3. 如果需要，进行变换
    if transform:
        target_image = transform(target_image)
    
    # 4. 确保模型在目标设备上运行
    model.to(device)
    
    # 5. 开启模型的评估模式和推理模式
    model.eval()
    with torch.inference_mode():
        # 在图像上添加一个额外的维度
        target_image = target_image.unsqueeze(dim=0)
    
        # 对带有额外维度的图像进行预测，并将结果发送到目标设备
        target_image_pred = model(target_image.to(device))
        
    # 6. 将预测的 logits 转换为预测概率（使用 torch.softmax() 进行多类别分类）
    target_image_pred_probs = torch.softmax(target_image_pred, dim=1)

    # 7. 将预测概率转换为预测标签
    target_image_pred_label = torch.argmax(target_image_pred_probs, dim=1)
    
    # 8. 绘制图像以及预测结果和预测概率
    plt.imshow(target_image.squeeze().permute(1, 2, 0)) # 确保图像大小适用于 matplotlib
    if class_names:
        title = f"预测: {class_names[target_image_pred_label.cpu()]} | 概率: {target_image_pred_probs.max().cpu():.3f}"
    else: 
        title = f"预测: {target_image_pred_label} | 概率: {target_image_pred_probs.max().cpu():.3f}"
    plt.title(title)
    plt.axis(False);

# Pred on our custom image
pred_and_plot_image(model=model_1,
                    image_path=custom_image_path,
                    class_names=class_names,
                    transform=custom_image_transform,
                    device=device)