1、处理数据样本的代码可能会变得混乱,难以维护。
2、为实现更好的可读性和模块化,希望处理数据集代码与模型训练代码实现脱钩与分离。
3、PyTorch提供了两个数据基元
torch.utils.data.DataLoader和torch.utils.data.Dataset。
它们允许你像使用自己的数据一样使用预加载的数据集。
(1)Dataset存储样本和它们对应的标签。
(2)DataLoader在Dataset周围包装一个可迭代对象,以便方便地访问样本。
以下是如何从TorchVision加载Fashion MNIST数据集的示例。Fashion MNIST是Zalando文章图像的数据集,由60000个训练示例和10000个测试示例组成。每个示例包括28×28灰度图像和来自10个类别之一的相关标签。
我们使用以下参数加载FashionMNIST数据集:
root 存储训练/测试数据的路径
train 定义了训练和测试dataset
download=True 如果根目录不可用,则从internet下载数据。
transform和targettransform指定特性和标签转换
代码如下
import torch
from torch.utils.data import Dataset
from torchvision import datasets
from torchvision.transforms import ToTensor
import matplotlib.pyplot as plt
training_data = datasets.FashionMNIST(
root="data",
train=True,
download=True,
transform=ToTensor()
)
test_data = datasets.FashionMNIST(
root="data",
train=False,
download=True,
transform=ToTensor()
)
我们可以像列表一样手动索引数据集:training_data[index]。我们使用matplotlib来可视化训练数据中的一些样本。
labels_map = {
0: "T-Shirt",
1: "Trouser",
2: "Pullover",
3: "Dress",
4: "Coat",
5: "Sandal",
6: "Shirt",
7: "Sneaker",
8: "Bag",
9: "Ankle Boot",
}
figure = plt.figure(figsize=(8, 8))
cols, rows = 3, 3
for i in range(1, cols * rows + 1):
sample_idx = torch.randint(len(training_data), size=(1,)).item()
img, label = training_data[sample_idx]
figure.add_subplot(rows, cols, i)
plt.title(labels_map[label])
plt.axis("off")
plt.imshow(img.squeeze(), cmap="gray")
plt.show()
自定义数据集类必须实现三个函数:init、len__和__getitem。
看看这个实现,FashionMNIST图像存储在目录img_dir中,其标签分别存储在CSV文件annotations_file中。
import os
import pandas as pd
from torchvision.io import read_image
class CustomImageDataset(Dataset):
def __init__(self, annotations_file, img_dir, transform=None, target_transform=None):
self.img_labels = pd.read_csv(annotations_file)
self.img_dir = img_dir
self.transform = transform
self.target_transform = target_transform
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 = read_image(img_path)
label = self.img_labels.iloc[idx, 1]
if self.transform:
image = self.transform(image)
if self.target_transform:
label = self.target_transform(label)
return image, label
__init__函数在实例化Dataset对象时运行一次。我们初始化包含图像、注释文件和两个转换的目录。
labels.csv内容
tshirt1.jpg, 0
tshirt2.jpg, 0
......
ankleboot999.jpg, 9
def __init__(self, annotations_file, img_dir, transform=None, target_transform=None):
self.img_labels = pd.read_csv(annotations_file)
self.img_dir = img_dir
self.transform = transform
self.target_transform = target_transform
__len__函数返回数据集中的样本数。
def __len__(self):
return len(self.img_labels)
__getitem_函数从给定索引idx处的数据集中加载并返回一个样本。基于索引,它识别图像在磁盘上的位置,
使用read_image将其转换为张量,从self.img_labels中的csv数据中检索相应的标签,调用其上的转换函数(如果适用),并返回元组中的张量图像和相应标签。
def __getitem__(self, idx):
img_path = os.path.join(self.img_dir, self.img_labels.iloc[idx, 0])
image = read_image(img_path)
label = self.img_labels.iloc[idx, 1]
if self.transform:
image = self.transform(image)
if self.target_transform:
label = self.target_transform(label)
return image, label
为了在PyTorch中定义神经网络,我们创建了一个继承自nn.Module的类。我们在__init__函数中定义了网络的层,并在forward函数中指定数据将如何通过网络。为了加速神经网络中的操作,我们将其移动到GPU(如果可用)。
# Get cpu or gpu device for training.
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using {device} device")
# Define model
class NeuralNetwork(nn.Module):
def __init__(self):
super().__init__()
self.flatten = nn.Flatten()
self.linear_relu_stack = nn.Sequential(
nn.Linear(28*28, 512),
nn.ReLU(),
nn.Linear(512, 512),
nn.ReLU(),
nn.Linear(512, 10)
)
def forward(self, x):
x = self.flatten(x)
logits = self.linear_relu_stack(x)
return logits
model = NeuralNetwork().to(device)
print(model)
为了训练模型,我们需要一个损失函数和一个优化器。
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
在单个训练循环中,模型在训练数据集上进行预测(分批提供给它),并反向预测预测误差以调整模型的参数。
def train(dataloader, model, loss_fn, optimizer):
size = len(dataloader.dataset)
model.train()
for batch, (X, y) in enumerate(dataloader):
X, y = X.to(device), y.to(device)
# Compute prediction error
pred = model(X)
loss = loss_fn(pred, y)
# Backpropagation
optimizer.zero_grad()
loss.backward()
optimizer.step()
if batch % 100 == 0:
loss, current = loss.item(), batch * len(X)
print(f"loss: {loss:>7f} [{current:>5d}/{size:>5d}]")