PyTorch 自定义数据集
准备数据
准备 COCO128 数据集,其是 COCO train2017 前 128 个数据。按 YOLOv5 组织的目录:
$ tree ~/datasets/coco128 -L 2
/home/john/datasets/coco128
├── images
│ └── train2017
│ ├── ...
│ └── 000000000650.jpg
├── labels
│ └── train2017
│ ├── ...
│ └── 000000000650.txt
├── LICENSE
└── README.txt
定义 Dataset
torch.utils.data.Dataset
是一个数据集的抽象类。自定义数据集时,需继承 Dataset
并覆盖如下方法:
__len__
:len(dataset)
获取数据集大小。__getitem__
:dataset[i]
访问第i
个数据。
详见:
自定义实现 YOLOv5 数据集的例子:
import os
from pathlib import Path
from typing import Any, Callable, Optional, Tuple
import numpy as np
import torch
import torchvision
from PIL import Image
class YOLOv5(torchvision.datasets.vision.VisionDataset):
def __init__(
self,
root: str,
name: str,
transform: Optional[Callable] = None,
target_transform: Optional[Callable] = None,
transforms: Optional[Callable] = None,
) -> None:
super(YOLOv5, self).__init__(root, transforms, transform, target_transform)
images_dir = Path(root) / 'images' / name
labels_dir = Path(root) / 'labels' / name
self.images = [n for n in images_dir.iterdir()]
self.labels = []
for image in self.images:
base, _ = os.path.splitext(os.path.basename(image))
label = labels_dir / f'{base}.txt'
self.labels.append(label if label.exists() else None)
def __getitem__(self, idx: int) -> Tuple[Any, Any]:
img = Image.open(self.images[idx]).convert('RGB')
label_file = self.labels[idx]
if label_file is not None: # found
with open(label_file, 'r') as f:
labels = [x.split() for x in f.read().strip().splitlines()]
labels = np.array(labels, dtype=np.float32)
else: # missing
labels = np.zeros((0, 5), dtype=np.float32)
boxes = []
classes = []
for label in labels:
x, y, w, h = label[1:]
boxes.append([
(x - w/2) * img.width,
(y - h/2) * img.height,
(x + w/2) * img.width,
(y + h/2) * img.height])
classes.append(label[0])
target = {}
target["boxes"] = torch.as_tensor(boxes, dtype=torch.float32)
target["labels"] = torch.as_tensor(classes, dtype=torch.int64)
if self.transforms is not None:
img, target = self.transforms(img, target)
return img, target
def __len__(self) -> int:
return len(self.images)
以上实现,继承了 VisionDataset
子类。其 __getitem__
返回了:
- image: PIL Image, 大小为
(H, W)
- target:
dict
, 含以下字段:boxes
(FloatTensor[N, 4]
): 真实标注框[x1, y1, x2, y2]
,x
范围[0,W]
,y
范围[0,H]
labels
(Int64Tensor[N]
): 上述标注框的类别标识
读取 Dataset
dataset = YOLOv5(Path.home() / 'datasets/coco128', 'train2017')
print(f'dataset: {len(dataset)}')
print(f'dataset[0]: {dataset[0]}')
输出:
dataset: 128
dataset[0]: (<PIL.Image.Image image mode=RGB size=640x480 at 0x7F6F9464ADF0>, {'boxes': tensor([[249.7296, 200.5402, 460.5399, 249.1901],
[448.1702, 363.7198, 471.1501, 406.2300],
...
[ 0.0000, 188.8901, 172.6400, 280.9003]]), 'labels': tensor([44, 51, 51, 51, 51, 44, 44, 44, 44, 44, 45, 45, 45, 45, 45, 45, 45, 45,
45, 50, 50, 50, 51, 51, 60, 42, 44, 45, 45, 45, 50, 51, 51, 51, 51, 51,
51, 44, 50, 50, 50, 45])})
预览:
使用 DataLoader
训练需要批量提取数据,可以使用 DataLoader :
dataset = YOLOv5(Path.home() / 'datasets/coco128', 'train2017',
transform=torchvision.transforms.Compose([
torchvision.transforms.ToTensor()
]))
dataloader = DataLoader(dataset, batch_size=64, shuffle=True,
collate_fn=lambda batch: tuple(zip(*batch)))
for batch_i, (images, targets) in enumerate(dataloader):
print(f'batch {batch_i}, images {len(images)}, targets {len(targets)}')
print(f' images[0]: shape={images[0].shape}')
print(f' targets[0]: {targets[0]}')
输出:
batch 0, images 64, targets 64
images[0]: shape=torch.Size([3, 480, 640])
targets[0]: {'boxes': tensor([[249.7296, 200.5402, 460.5399, 249.1901],
[448.1702, 363.7198, 471.1501, 406.2300],
...
[ 0.0000, 188.8901, 172.6400, 280.9003]]), 'labels': tensor([44, 51, 51, 51, 51, 44, 44, 44, 44, 44, 45, 45, 45, 45, 45, 45, 45, 45,
45, 50, 50, 50, 51, 51, 60, 42, 44, 45, 45, 45, 50, 51, 51, 51, 51, 51,
51, 44, 50, 50, 50, 45])}
batch 1, images 64, targets 64
images[0]: shape=torch.Size([3, 248, 640])
targets[0]: {'boxes': tensor([[337.9299, 167.8500, 378.6999, 191.3100],
[383.5398, 148.4501, 452.6598, 191.4701],
[467.9299, 149.9001, 540.8099, 193.2401],
[196.3898, 142.7200, 271.6896, 190.0999],
[134.3901, 154.5799, 193.9299, 189.1699],
[ 89.5299, 162.1901, 124.3798, 188.3301],
[ 1.6701, 154.9299, 56.8400, 188.3700]]), 'labels': tensor([20, 20, 20, 20, 20, 20, 20])}
源码
参考
APIs:
GoCoding 个人实践的经验分享,可关注公众号!
PyTorch 自定义数据集的更多相关文章
- [转载]pytorch自定义数据集
为什么要定义Datasets: PyTorch提供了一个工具函数torch.utils.data.DataLoader.通过这个类,我们在准备mini-batch的时候可以多线程并行处理,这样可以加快 ...
- Pytorch划分数据集的方法
之前用过sklearn提供的划分数据集的函数,觉得超级方便.但是在使用TensorFlow和Pytorch的时候一直找不到类似的功能,之前搜索的关键字都是"pytorch split dat ...
- pytorch加载语音类自定义数据集
pytorch对一下常用的公开数据集有很方便的API接口,但是当我们需要使用自己的数据集训练神经网络时,就需要自定义数据集,在pytorch中,提供了一些类,方便我们定义自己的数据集合 torch.u ...
- MMDetection 快速开始,训练自定义数据集
本文将快速引导使用 MMDetection ,记录了实践中需注意的一些问题. 环境准备 基础环境 Nvidia 显卡的主机 Ubuntu 18.04 系统安装,可见 制作 USB 启动盘,及系统安装 ...
- Scaled-YOLOv4 快速开始,训练自定义数据集
代码: https://github.com/ikuokuo/start-scaled-yolov4 Scaled-YOLOv4 代码: https://github.com/WongKinYiu/S ...
- torch_13_自定义数据集实战
1.将图片的路径和标签写入csv文件并实现读取 # 创建一个文件,包含image,存放方式:label pokemeon\\mew\\0001.jpg,0 def load_csv(self,file ...
- Tensorflow2 自定义数据集图片完成图片分类任务
对于自定义数据集的图片任务,通用流程一般分为以下几个步骤: Load data Train-Val-Test Build model Transfer Learning 其中大部分精力会花在数据的准备 ...
- Pytorch自定义数据库
1)前言 虽然torchvision.datasets中已经封装了好多通用的数据集,但是我们在使用Pytorch做深度学习任务的时候,会面临着自定义数据库来满足自己的任务需要.如我们要训练一个人脸关键 ...
- [炼丹术]YOLOv5训练自定义数据集
YOLOv5训练自定义数据 一.开始之前的准备工作 克隆 repo 并在Python>=3.6.0环境中安装requirements.txt,包括PyTorch>=1.7.模型和数据集会从 ...
随机推荐
- php之PDO连接mysql数据库,增删改查等等操作实例
我们使用传统的 mysql_connect .mysql_query方法来连接查询数据库时,如果过滤不严就有SQL注入风险,导致网站被攻击. 虽然可以用mysql_real_escape_string ...
- 一个简单且易上手的 Spring boot 后台管理框架-->EL-ADMIN
一个简单且易上手的 Spring boot 后台管理框架 后台源码 前台源码
- Hyperbase数据迁移
原老集群有100台服务器,新增90台服务器和原来的服务器构成新Hyperbase集群最初考虑有两种方案distcp和snapshot,由于distcp进行数据迁移时不在HBase本身控制范围内,故选用 ...
- 【poj 1961】Period(字符串--KMP 模版题)
题意:给你一个字符串,求这个字符串到第 i 个字符为止的重复子串的个数. 解法:判断重复子串的语句很重要!!if (p && i%(i-p)==0) printf("%d % ...
- 2019牛客多校 Round1
Solved:4 Rank:143 A Equivalent Prefixes 题意:求一个最大的r满足在A,B两个数组中1,r里所有的子区间RMQ相等 题解:单调队列秒 #include <b ...
- c语言实现--带头结点单链表操作
可能是顺序表研究的细致了一点,单链表操作一下子就实现了.这里先实现带头结点的单链表操作. 大概有以下知识点. 1;结点:结点就是单链表中研究的数据元素,结点中存储数据的部分称为数据域,存储直接后继地址 ...
- Codeforces Round #555 (Div. 3) E. Minimum Array (贪心,二分,set)
题意:给你两个长度为\(n\)的数组\(a\)和\(b\),元素值在\([0,n-1]\),可以对\(b\)数组的元素任意排序,求新数组\(c\),满足\(c_i=(a_i+b_i)\ mod\ n\ ...
- Smith Numbers POJ - 1142 暴力递归枚举
题意: 给你一个数x,把这个分解成素数之积(假设是x1*x2*x3),如果 x的每一数位的和 等于 x1每一数位的和加上x2每一数位的和加上x3每一数位的和,那么他就是题目要找的数 示例: ...
- Codeforces Round #653 (Div. 3) D. Zero Remainder Array (数学,模拟)
题意:有一组数,刚开始时\(x=0\),每次可以让\(x\)++或让某一个元素+=\(x\)后\(x\)++,每个元素只能加一次\(x\),问最少操作多少次使得所有元素能被\(k\)整除. 题解:每个 ...
- 微服务架构学习Day01-SpringBoot入门
基本概念 SpringBoot的优点: 可以创建独立的Spring应用 SpringBoot嵌入Tomcat,Jetty和Unsertow, 不需要部署war文件 根据需要通过maven获取start ...