PyTorch的dataloader制作自定义数据集

PyTorch的dataloader是用于读取训练数据的工具,它可以自动将数据分割成小batch,并在训练过程中进行数据预处理。以下是制作PyTorch的dataloader的简单步骤:

  1. 导入必要的库

    import torch
    from torch.utils.data import DataLoader, Dataset

  2. 定义数据集类 需要自定义一个继承自torch.utils.data.Dataset的类,在该类中实现__len____getitem__方法。

    class MyDataset(Dataset):
    def init(self, data):
    self.data = data

    复制代码
     def __len__(self):
         return len(self.data)
     
     def __getitem__(self, index):
         # 返回第index个数据样本
         return self.data[index]
  3. 创建数据集实例

    data = [1, 2, 3, 4, 5]
    dataset = MyDataset(data)

  4. 创建dataloader实例

使用torch.utils.data.DataLoader创建dataloader实例,可以设置batch_sizeshuffle等参数。

复制代码
dataloader = DataLoader(dataset, batch_size=2, shuffle=True)
  1. 使用dataloader读取数据

    for batch in dataloader:
    # batch为一个batch的数据,可以直接用于训练
    print(batch)

以上是制作PyTorch的dataloader的简单步骤,根据实际需求可以进行更复杂的操作,如数据增强、并行读取等。

5.已经分类的文件生成标注文件

假设你已经将所有的图片按照类别分别放到了十个文件夹中,可以使用以下代码生成标注文件:

复制代码
import os
# 定义图片所在的文件夹路径和标注文件的路径
img_dir = '/path/to/image/directory'
ann_file = '/path/to/annotation/file.txt'
# 遍历每个类别文件夹中的图片,将标注信息写入到标注文件中
with open(ann_file, 'w') as f:
    for class_id in range(1, 11):
        class_dir = os.path.join(img_dir, 'class{}'.format(class_id))
        for filename in os.listdir(class_dir):
            if filename.endswith('.jpg'):
                # 写入图片的文件名和类别
                f.write('{} {}\n'.format(filename, class_id))
相关推荐
code_pgf4 分钟前
Jetson Orin NX 16G设备上配置AI服务自动启动的方案,包括Ollama、llama-server和OpenClaw Gateway三个组件
数据库·人工智能·安全·gateway·边缘计算·llama
前端付豪6 分钟前
实现 AI 回复支持 Markdown 渲染
前端·人工智能·markdown
数智大号6 分钟前
北京InfoComm China 2026高峰会议紧密契合国家“十五五”规划,人工智能引领科技未
人工智能
码路飞7 分钟前
Mistral Small 4 上手实测:119B 参数只激活 6B,开源模型卷到这地步了?
人工智能·llm
用户4815930195917 分钟前
为什么 Claude Code 值得深入学
人工智能
七夜zippoe8 分钟前
边缘计算:Python在IoT中的应用
python·物联网·esp32·边缘计算·iot
2301_818419019 分钟前
Python虚拟环境(venv)完全指南:隔离项目依赖
jvm·数据库·python
蒸汽求职9 分钟前
蒸汽教育求职分享:2026年数据工程师就业优势分析与职业发展路径指南
数据库·人工智能·面试·求职招聘·美国求职
火山引擎开发者社区10 分钟前
重磅发布|Scale‑SWE 构造 10 万级真实 SWE 数据集,火山引擎沙箱底座重塑代码智能体训练
人工智能
MgArcher10 分钟前
python基础:使用list和tuple
python