05-PyTorch自定义数据集

本文为PyTorch 自定义数据集的学习笔记,对原文进行了翻译和编辑,本系列课程介绍和目录在《使用 PyTorch 进行深度学习系列》课程介绍

文章将最先在我的博客发布,其他平台因为限制不能实时修改。

在微信公众号内无法嵌入超链接,可以点击底部阅读原文获得更好的阅读体验。


对于机器学习中的许多不同问题,我们采取的步骤都是相似的。PyTorch 有许多内置数据集,用于大量机器学习基准测试。除此之外也可以自定义数据集,本问将使用我们自己的披萨、牛排和寿司图像数据集,而不是使用内置的 PyTorch 数据集。具体来说,我们将使用 torchvision.datasets 以及我们自己的自定义 Dataset 类来加载食物图像,然后我们将构建一个 PyTorch 计算机视觉模型,希望对三种物体进行分类。

building a pipeline to load in food images and then building a pytorch model to classify those food images

什么是自定义数据集?

自定义数据集是与您正在处理的特定问题相关的数据集合。本质上,自定义数据集几乎可以由任何内容组成。

例如,如果我们正在构建像 Nutrify 这样的食物图像分类应用程序,我们的自定义数据集可能是食物图像。

或者,如果我们尝试构建一个模型来对网站上基于文本的评论是正面还是负面进行分类,那么我们的自定义数据集可能是现有客户评论及其评级的示例。

或者,如果我们尝试构建声音分类应用程序,我们的自定义数据集可能是声音样本及其样本标签。

或者,如果我们试图为在我们网站上购买商品的客户构建推荐系统,我们的自定义数据集可能是其他人购买过的产品的示例。

different pytorch domain libraries can be used for specific PyTorch problems

PyTorch 包含许多现有函数,可加载 TorchVision, TorchText, TorchAudioTorchRec 库中的各种自定义数据集。

在这种情况下,我们总是可以子类化 torch.utils.data.Dataset 并根据我们的喜好自定义它。

0.导入 PyTorch

import torch
from torch import nn

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

1. 获取数据

我们将使用的数据是 Food101 dataset 的子集。

Food101 是流行的计算机视觉基准,因为它包含 101 种不同食物的 1000 张图像,总共 101,000 张图像(75,750 个训练图像和 25,250 个测试图像)。

如果您想查看数据的来源,可以查看以下资源:

  1. 原始 Food101 数据集和论文网站。
  2. torchvision.datasets.Food101 - 我为此笔记本下载的数据版本。
  3. extras/04_custom_data_creation.ipynb -格式化 Food101 数据集以用于此笔记本的笔记本。
  4. data/pizza_steak_sushi.zip - 来自 Food101 的披萨、牛排和寿司图像的 zip 存档,使用上面链接的笔记本创建。

使用上述第 4 点下载的数据,或者使用以下代码下载数据并解压:

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)

2. 数据准备

将数据文件夹整理为以下目录结构,例如, pizza 的所有图像都包含在 pizza/ 目录中。:

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

OUt:

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

2.1 可视化图像

随机选择一些图形并且可视化:

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

out:

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

img


我们可以使用 matplotlib.pyplot.imshow() 执行相同的操作,只不过我们必须首先将图像转换为 NumPy 数组。

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);

img

3. 转换数据

PyTorch 有几种不同类型的预构建数据集和数据集加载器,具体取决于您正在处理的问题。

Problem space 问题空间 Pre-built Datasets and Functions 预构建的数据集和函数
**Vision ** torchvision.datasets
**Audio ** torchaudio.datasets
**Text ** torchtext.datasets
**Recommendation system ** torchrec.datasets

由于我们正在处理视觉问题,因此我们将使用 torchvision.datasets 来获取数据加载功能,并使用 torchvision.transforms 来转换数据。

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
])

除了 transforms.Resize()transforms.RandomHorizontalFlip()transforms.ToTensor()torchvision.transforms.Compose(),在 PyTorch torchvision.transforms 文档介绍了更多函数。

测试一下转换函数:

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)

4. 方式一:使用 ImageFolder 加载图像数据

由于我们的数据采用标准图像分类格式,因此我们可以使用类 torchvision.datasets.ImageFolder

我们可以向它传递目标图像目录的文件路径以及我们想要对图像执行的一系列转换。

让我们在数据文件夹 train_dirtest_dir 上进行测试,传入 transform=data_transform 将图像转换为张量。

# 使用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}")

out:

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()
           )

dataset 的属性

#数据集的类别
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)

对 dataset 进行索引

我们可以对 train_datatest_dataDataset 进行索引来查找样本及其目标标签。

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'>

我们的图像现在采用张量的形式(形状 [3, 64, 64] ),标签采用与特定类相关的整数形式(使用class_to_idx 属性调用:{‘pizza’: 0, ‘steak’: 1, ‘sushi’: 2}) 。

我们使用 matplotlib 绘制单个图像?

首先必须进行排列(重新排列其维度的顺序)以使其兼容。因为我们的图像尺寸采用格式 CHW (颜色通道、高度、宽度),但 matplotlib 使用 HWC (高度、宽度、颜色通道)。

# 重新排列图像维度
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);

out:

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

Image

因为图像的大小从 512x512 调整为 64x64 像素,质量较差。

4.1 将加载的 Dataset 转为 DataLoader

我们将使用 torch.utils.data.DataLoaderDataset 转换为 DataLoader 使它们可迭代,以便模型可以学习样本和目标(特征和标签)之间的关系。

为了简单起见,我们将使用 batch_size=1num_workers=1

# 将加载的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

out:

(<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}")

out:

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

我们现在可以使用这些 DataLoader 进行训练和测试循环来训练模型。在此之前,让我们看看另一种自定义加载图像方式:

5.方法二:使用自定义 Dataset 加载图像数据

如果像 torchvision.datasets.ImageFolder() 这样的预构建 Dataset 创建器不存在怎么办?例如,我们的标签和图像储存 csv 文件中,同时图片文件夹又没有按照标准方式进行储存,我们则可以使用自定义数据集。

导入相关库:

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

5.1 创建函数来获取类名

还记得 torchvision.datasets.ImageFolder() 实例如何允许我们使用 classesclass_to_idx 属性吗?

此实例定义了classesclass_to_idx属性,为了方便我们定义一个函数来根据文件所在的文件夹的名称来定义类(你也可以使用其他方式,比如从 csv 中读取相应文件的类名)。

# 创建函数以在目标目录中查找类别
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() 函数。

find_classes(train_dir)

out:

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

5.2 创建自定义 Dataset 来复制 ImageFolder

接下来是一段相当长的代码……

# 编写自定义数据集类(继承自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. 创建一个 ImageFolderCustom 类,并且继承 torch.utils.data.Dataset
  2. 使用 targ_dir 参数(目标数据目录)和 transform 参数初始化我们的子类(因此我们可以选择在需要时转换数据)。transform 可以不填,后文会定义一个transform 函数
  3. paths (目标图像的路径)、 transform (我们可能想要使用的转换,可以是 None )、 classesclass_to_idx (来自我们的 find_classes() 函数)。
  4. 创建一个函数来从文件加载图像并返回它们,这可以使用 PILtorchvision.io (用于视觉数据的输入/输出)。
  5. 覆盖 torch.utils.data.Dataset__len__ 方法以返回 Dataset 中的样本数,建议但不是必需的。这样您就可以调用 len(Dataset)
  6. 覆盖 torch.utils.data.Dataset__getitem__ 方法以从 Dataset 返回单个样本,这是必需的。

5.3 实例化 Dataset

# 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

out:

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

5.4 检查 Dataset 属性

让我们尝试在新的 Dataset 上调用 len() 并找到 classesclass_to_idx 属性。

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}

5.4 将自定义加载的图像转换为 DataLoader 对象

我们可以使用与之前非常相似的步骤,只不过这次我们将使用自定义创建的 Dataset

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)

6. 数据增强 data augmentation

可以在 torchvision.transforms 文档中查看它们。

变换的目的是以某种方式改变你的图像。除了将图像变成张量,还有裁剪、随机擦除、随机水平镜像等一部分或随机旋转它们。进行这种转换通常称为数据增强。数据增强是通过人为增加训练集多样性的方式更改数据的过程。

您可以在 PyTorch 的变换示例中看到使用 torchvision.transforms 在图像上执行数据增强的许多不同示例。

机器学习就是利用随机性的力量,研究表明随机变换(如 transforms.RandAugment()transforms.TrivialAugmentWide())通常比手工选择的变换表现更好。

04-trivial-augment-being-using-in-PyTorch-resize

TrivialAugment 是最近对各种 PyTorch 视觉模型进行最先进的训练升级时使用的成分之一。

transforms.TrivialAugmentWide() 中需要注意的主要参数是 num_magnitude_bins=31

它定义了将选择多少范围的强度值来应用特定变换, 0 是无范围, 31 是最大范围(最高强度的最高机会)。

我们可以将 transforms.TrivialAugmentWide() 合并到 transforms.Compose() 中。

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


7.模型 0:没有数据增强的 TinyVGG

我们已经了解了如何将数据从文件夹中的图像转换为张量。现在让我们构建一个计算机视觉模型,看看我们是否可以对图像进行分类:披萨、牛排还是寿司。

首先,我们将从一个简单的转换开始,仅将图像大小调整为 (64, 64) 并将它们转换为张量。

7.1 为模型 0 创建转换并加载数据

# 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

7.2 创建 TinyVGG 模型类

让我们重新创建相同的模型,只不过这次我们将使用彩色图像而不是灰度图像(对于 RGB 像素,使用 in_channels=3 而不是 in_channels=1 )。

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)
  )
)

7.3 使用 torchinfo 了解模型中每一层形状的变化

使用 print(model) 打印模型可以让我们了解模型的情况。然而,从模型中获取信息的令一个有用方法是使用 torchinfotorchinfo 附带一个 summary() 方法,该方法采用 PyTorch 模型以及 input_shape 并返回张量在模型中移动时发生的情况。

安装:

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

model_0 info

Total params ,我们模型中的参数总数, Estimated Total Size (MB) 是我们模型的大小。

7.4 创建训练和测试循环函数

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

现在让我们对 test_step() 函数执行相同的操作。

7.6 创建一个 train() 函数来组合 train_step() 和 test_step()

现在我们需要一种方法将 train_step()test_step() 函数组合在一起。

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

7.7 训练和评估模型 0

# 设置随机种子
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} 秒")

out:

  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

7.8 绘制模型 0 的损失曲线

# 检查 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()

这段代码定义了一个名为 plot_loss_curves 的函数,用于绘制训练曲线。函数接受一个结果字典 (results) 作为输入。

在函数内部,首先从结果字典中获取训练损失值 (loss)、测试损失值 (test_loss)、训练准确率值 (accuracy) 和测试准确率值 (test_accuracy)。

然后,确定有多少个 epoch,使用 range 函数生成一个表示 epoch 的范围。

接下来,设置绘图的大小,并创建一个包含两个子图的图表。

第一个子图绘制损失曲线,使用 plt.plot 函数绘制训练损失值和测试损失值随 epoch 变化的曲线。设置标题为 “损失”,x 轴标签为 “Epochs”,并添加图例。

第二个子图绘制准确率曲线,使用 plt.plot 函数绘制训练准确率值和测试准确率值随 epoch 变化的曲线。设置标题为 “准确率”,x 轴标签为 “Epochs”,并添加图例。

loss_curves

8. 理想的损失曲线应该是什么样的?

查看训练和测试损失曲线是查看模型是否过度拟合的好方法。过度拟合模型是一种在训练集上比在验证/测试集上表现更好(通常有相当大的优势)的模型。

如果您的训练损失远低于测试损失,则您的模型过度拟合。比如说,它在训练中学习的模式太好了,而这些模式并没有推广到测试数据。

另一方面是当你的训练和测试损失没有你想要的那么低时,这被认为是欠拟合。

训练和测试损失曲线的理想位置是它们彼此紧密对齐。

different training and test loss curves illustrating overfitting, underfitting and the ideal loss curves

  • 左:如果您的训练和测试损失曲线没有您想要的那么低,则被认为是欠拟合。
  • 中:当您的测试/验证损失高于训练损失时,这被认为是过度拟合。
  • 右图:理想的情况是训练和测试损失曲线随着时间的推移保持一致。这意味着您的模型具有良好的泛化能力。损失曲线可以做更多的组合和不同的事情,有关这些的更多信息,请参阅 Google 的解释损失曲线指南

8.1 如何处理过度拟合

防止过度拟合的常用技术称为正则化

让我们讨论一些防止过度拟合的方法。

防止过拟合的方法 说明
获取更多数据 拥有更多数据使模型有更多机会学习模式,这些模式可能更适用于新示例。
简化您的模型 如果当前模型已经过度拟合训练数据,则模型可能过于复杂。这意味着它对数据模式的学习太好了,无法很好地泛化到未见过的数据。简化模型的一种方法是减少其使用的层数或减少每层中隐藏单元的数量。
使用数据增强 数据增强以某种方式操纵训练数据,使模型更难学习,因为它人为地为数据添加了更多多样性。如果模型能够学习增强数据中的模式,则该模型可能能够更好地泛化到未见过的数据。
使用迁移学习 迁移学习涉及利用模型已学会的模式(也称为预训练权重)作为您自己的任务的基础。在我们的例子中,我们可以使用一种在多种图像上进行预训练的计算机视觉模型,然后稍微调整它以更专门针对食物图像。
使用 dropout 层 Dropout 层随机删除神经网络中隐藏层之间的连接,有效地简化了模型,同时也使剩余的连接变得更好。有关更多信息,请参阅 torch.nn.Dropout()
使用学习率衰减 这里的想法是在模型训练时慢慢降低学习率。这类似于伸手去拿沙发后面的硬币。距离越近,脚步就越小。与学习率相同,越接近收敛 convergence,您希望权重更新越小。
Use early stopping 使用提前停止 Early stopping stops model training before it begins to overfit. As in, say the model’s loss has stopped decreasing for the past 10 epochs (this number is arbitrary), you may want to stop the model training here and go with the model weights that had the lowest loss (10 epochs prior). 提前停止会在模型开始过度拟合之前停止训练。例如,假设模型的损失在过去 10 个时期内已停止减少(该数字是任意的),您可能希望在此处停止模型训练并使用损失最低的模型权重(之前的 10 个时期)。

当您开始构建越来越多的深度模型时,您会发现由于深度学习非常擅长学习数据模式,因此处理过度拟合是深度学习的主要问题之一

8.2 如何处理欠拟合

防止欠拟合的方法 说明
向模型添加更多层/单元 如果您的模型拟合不足,它可能没有足够的能力来学习预测所需的数据模式/权重/表示。为模型添加更多预测能力的一种方法是增加这些层中隐藏层/单元的数量。
调整学习率 也许你的模型的学习率一开始就太高了。而且它在每个时期都试图过多地更新权重,结果却没有学到任何东西。在这种情况下,您可以降低学习率并看看会发生什么。
Use transfer learning 使用迁移学习 迁移学习能够防止过度拟合和欠拟合。它涉及使用以前工作模型中的模式并根据您自己的问题进行调整。
Train for longer 训练时间更长 有时模型只是需要更多时间来学习数据的表示。如果您发现在较小的实验中您的模型没有学到任何东西,也许让它训练更多的时期可能会带来更好的性能。
使用较少的正则化 也许您的模型拟合不足,因为您试图防止过度拟合。抑制正则化技术可以帮助您的模型更好地拟合数据。

8.3 过拟合和欠拟合之间的平衡

防止过度拟合和欠拟合可能是机器学习研究最活跃的领域。

迁移学习不是手动设计不同的过拟合和欠拟合技术,而是使您能够在与您的问题空间类似的问题空间中采用已经训练的模型(例如 Hugging Face 的模型)并将其应用到您自己的数据集。

我们将在稍后的笔记本中看到迁移学习的力量。

9. 模型 1:具有数据增强功能的 TinyVGG

这次,让我们加载数据并使用数据增强来看看它是否能改善我们的结果。

我们将编写一个训练转换以包含 transforms.TrivialAugmentWide() 以及调整图像大小并将图像转换为张量。我们将对测试转换执行相同的操作,但不进行数据增强。

9.1 使用数据增强创建转换

# 使用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()
])

9.2 创建训练和测试 Dataset 和 DataLoader

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)

9.3 构建和训练模型 1

# 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

9.4 绘制模型 1 的损失曲线

plot_loss_curves(model_1_results)

10. 比较模型结果

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

image-20230929234427353

现在我们可以使用 matplotlib 编写一些绘图代码来一起可视化 model_0model_1 的结果。

# 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();

看起来我们的模型表现同样糟糕并且有点波动(指标急剧上升和下降)。

11.使用模型进行预测

我们预测这一张图来验证模型:

下载 (22)

# 下载自定义图像
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')

好吧,这些仍然是 logit 形式(模型的原始输出称为 logits)。让我们将它们从 logits -> 预测概率 -> 预测标签转换:

# 打印预测的 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')

但当然我们的预测标签仍然是索引/张量的形式。我们可以通过在 class_names 列表上建立索引将其转换为字符串类名预测。

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

11.3 将以上预测放在一起:构建函数

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)

再次竖起两个大拇指!


额外资料:

PyTorch 和深度学习的三大错误:

  1. 错误的数据类型 - 当您的数据为 torch.uint8 时,您的模型期望 torch.float32
  2. 错误的数据形状 - 当您的数据为 [color_channels, height, width] 时,您的模型预期为 [batch_size, color_channels, height, width]
  3. 错误的设备 - 您的模型位于 GPU 上,但您的数据位于 CPU 上。

阅读资料

感谢

感谢原作者 Daniel Bourke,访问https://www.learnpytorch.io/可以阅读英文原文,点击原作者的 Github 仓库:https://github.com/mrdbourke/pytorch-deep-learning/可以获得帮助和其他信息。

本文同样遵守遵守 MIT license,不受任何限制,包括但不限于权利

使用、复制、修改、合并、发布、分发、再许可和/或出售。但需标明原始作者的许可信息:renhai-lab:https://blog.renhai.online/


05-PyTorch自定义数据集
https://blog.renhai.online/archives/DL-05-pytorch-custom-datasets
作者
Renhai
发布于
2023年10月09日
更新于
2024年06月17日
许可协议