PyTorch 自定义数据集示例教程
项目目录结构及介绍
pytorch-custom-dataset-examples/
├── data/
│ └── ... (数据文件)
├── src/
│ ├── custom_dataset.py
│ └── ... (其他源代码文件)
├── .gitignore
├── LICENSE
├── README.md
└── ... (其他配置文件)
data/
: 存储数据文件的目录。src/
: 包含项目的主要源代码文件,如custom_dataset.py
。.gitignore
: 指定 Git 忽略的文件和目录。LICENSE
: 项目的许可证文件。README.md
: 项目的说明文档。
项目的启动文件介绍
项目的启动文件通常是 src/custom_dataset.py
,该文件定义了自定义数据集类 CustomDatasetFromFile
,用于加载和处理数据。
from torch.utils.data import Dataset
import glob
from PIL import Image
import numpy as np
class CustomDatasetFromFile(Dataset):
def __init__(self, folder_path):
"""
A dataset example where the class is embedded in the file names.
This data example also does not use any torch transforms.
Args:
folder_path (string): path to image folder
"""
# Get image list
self.image_list = glob.glob(folder_path + '*')
# Calculate len
self.data_len = len(self.image_list)
def __getitem__(self, index):
# Get image name from the pandas df
single_image_path = self.image_list[index]
# Open image
im_as_im = Image.open(single_image_path)
# Do some operations on image
# Convert to numpy, dim = 28x28
im_as_np = np.asarray(im_as_im) / 255
# Add channel dimension, dim = 1x28x28
# Note: You do not need to do this if you are reading RGB images
# or if you add a transform that adds the channel dimension.
im_as_np = np.expand_dims(im_as_np, axis=0)
return im_as_np
def __len__(self):
return self.data_len
项目的配置文件介绍
项目中没有明确的配置文件,但可以通过修改 src/custom_dataset.py
中的参数来调整数据集的加载和处理方式。例如,可以修改 folder_path
参数来指定不同的数据目录。
folder_path = 'data/' # 修改此路径以指向不同的数据目录
dataset = CustomDatasetFromFile(folder_path)
通过这种方式,可以灵活地配置和使用自定义数据集。