step1:自定义数据集类必须实现三个函数:__init__、__len__和__getitem__。
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
"""
__init__函数在实例化数据集对象时运行一次。我们初始化 包含图像、注释文件和两个转换的目录
"""
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
"""
__getitem__函数加载并从给定索引处的数据集返回样本。 根据索引,它识别图像在磁盘上的位置,将
其转换为张量,使用 ,检索 中 CSV 数据的相应标签,调用其上的转换函数(如果适用),并返回 元
组中的张量图像和相应的标签。
"""
step2:准备数据以使用数据加载程序进行训练。(检索数据集的特征并一次标记一个样本。在训练模型时,我们通常希望 在“小批量”中传递样本,在每个时期重新洗牌数据以减少模型过度拟合,并使用 Python 的 加快数据检索速度。)
from torch.utils.data import DataLoader
train_dataloader = DataLoader(training_data, batch_size=64, shuffle=True)
#training_data为step1创建的
test_dataloader = DataLoader(test_data, batch_size=64, shuffle=True)
step3:遍历 DataLoader。(我们已将该数据集加载到DataLoader中,并可以根据需要循环访问该数据集。 下面的每次迭代都会返回一批 图像and标签(分别包含特征和标注)。 因为我们指定了shuffle=True(随机)
,在我们遍历所有批次后,数据将被洗牌。DataLoader
train_features
train_labels
batch_size=64
shuffle=True
)
# Display image and label.
train_features, train_labels = next(iter(train_dataloader))
print(f"Feature batch shape: {train_features.size()}")
print(f"Labels batch shape: {train_labels.size()}")
img = train_features[0].squeeze()
label = train_labels[0]
plt.imshow(img, cmap="gray")
plt.show()
print(f"Label: {label}")
结果: