torchvision.datasets.ImageFolder()使用详解
时间: 2024-09-17 07:09:26 浏览: 91
`torchvision.datasets.ImageFolder()`是PyTorch提供的一个非常方便的数据集类,用于处理图像数据,特别是那些按照文件夹结构组织的分类数据集,例如常见的ImageNet。它假设数据集根目录下有两类文件夹:一类包含训练图像(通常是类别的名称),另一类包含对应的标签索引文件。
使用`ImageFolder()`的步骤通常包括:
1. **初始化**:你需要提供数据集的路径作为参数,如`ImageFolder(root='path/to/dataset')`。
2. **加载数据**:创建`ImageFolder`实例后,你可以通过`data_loader = torch.utils.data.DataLoader(ImageFolder(root), batch_size=64, shuffle=True)`创建一个数据加载器(DataLoader),它会自动读取图片并将其转换成Tensor。
3. **遍历数据**:数据加载器会返回批次(batch)的样本,每个批次包含图片张数(batch_size)大小的一组样本,每张图片都附带相应的标签。
例如,你可以像这样查看和操作数据:
```python
dataset = ImageFolder('my_dataset')
for images, labels in dataset:
# images 是一个包含batch_size张图片的张量
# labels 是一个包含对应图片标签的整数列表
```
相关问题
pillow中ImageFolder()详解
Pillow是Python中一个非常强大的图像处理库,可以对图像进行各种操作。其中,ImageFolder()是Pillow中的一个函数,用于读取文件夹中的所有图像文件,将其加载为一个数据集。
具体来说,ImageFolder()函数可以接受两个参数:文件夹路径和一个可选的变换函数。文件夹路径指定了包含图像文件的文件夹,变换函数可以对图像进行一些转换,比如裁剪、缩放等。
ImageFolder()函数返回一个数据集,其中每个元素都是一个图像和其对应的标签。标签是从文件夹名字中提取的,所以需要满足一定的命名规范。
下面是一个使用ImageFolder()函数读取数据集的示例代码:
```
from torchvision import datasets, transforms
# 定义变换函数
transform = transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
# 加载数据集
dataset = datasets.ImageFolder('data', transform=transform)
# 获取数据集大小和类别数
print('Dataset size:', len(dataset))
print('Number of classes:', len(dataset.classes))
```
在上面的代码中,我们首先定义了一个变换函数,其中包含了随机裁剪、随机水平翻转、归一化等操作。然后使用ImageFolder()函数读取data文件夹中的图像文件,并应用变换函数,得到一个数据集。最后,我们输出了数据集的大小和类别数,分别是200张图像和2个类别。
yolov7train.py详解
yolov7train.py 是使用 YOLOv7 算法进行目标检测的训练脚本。下面对 yolov7train.py 的主要代码进行简单的解释:
1. 导入相关库
```python
import argparse
import yaml
import time
import torch
from torch.utils.data import DataLoader
from torchvision import datasets
from models.yolov7 import Model
from utils.datasets import ImageFolder
from utils.general import (
check_img_size, non_max_suppression, apply_classifier, scale_coords,
xyxy2xywh, plot_one_box, strip_optimizer, set_logging)
from utils.torch_utils import (
select_device, time_synchronized, load_classifier, model_info)
```
这里导入了 argparse 用于解析命令行参数,yaml 用于解析配置文件,time 用于记录时间,torch 用于神经网络训练,DataLoader 用于读取数据集,datasets 和 ImageFolder 用于加载数据集,Model 用于定义 YOLOv7 模型,各种工具函数用于辅助训练。
2. 定义命令行参数
```python
parser = argparse.ArgumentParser()
parser.add_argument('--data', type=str, default='data.yaml', help='dataset.yaml path')
parser.add_argument('--hyp', type=str, default='hyp.yaml', help='hyperparameters path')
parser.add_argument('--epochs', type=int, default=300)
parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs')
parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='[train, test] image sizes')
parser.add_argument('--rect', action='store_true', help='rectangular training')
parser.add_argument('--resume', nargs='?', const='yolov7.pt', default=False, help='resume most recent training')
parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')
parser.add_argument('--notest', action='store_true', help='only test final epoch')
parser.add_argument('--evolve', action='store_true', help='evolve hyperparameters')
parser.add_argument('--bucket', type=str, default='', help='gsutil bucket')
opt = parser.parse_args()
```
这里定义了许多命令行参数,包括数据集路径、超参数路径、训练轮数、批量大小、图片大小、是否使用矩形训练、是否从最近的检查点恢复训练、是否只保存最终的检查点、是否只测试最终的模型、是否进行超参数进化、gsutil 存储桶等。
3. 加载数据集
```python
with open(opt.data) as f:
data_dict = yaml.load(f, Loader=yaml.FullLoader)
train_path = data_dict['train']
test_path = data_dict['test']
num_classes = data_dict['nc']
names = data_dict['names']
train_dataset = ImageFolder(train_path, img_size=opt.img_size[0], rect=opt.rect)
test_dataset = ImageFolder(test_path, img_size=opt.img_size[1], rect=True)
batch_size = opt.batch_size
train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=8, pin_memory=True, collate_fn=train_dataset.collate_fn)
test_dataloader = DataLoader(test_dataset, batch_size=batch_size * 2, num_workers=8, pin_memory=True, collate_fn=test_dataset.collate_fn)
```
这里读取了数据集的配置文件,包括训练集、测试集、类别数和类别名称等信息。然后使用 ImageFolder 加载数据集,设置图片大小和是否使用矩形训练。最后使用 DataLoader 加载数据集,并设置批量大小、是否 shuffle、是否使用 pin_memory 等参数。
4. 定义 YOLOv7 模型
```python
model = Model(opt.hyp, num_classes, opt.img_size)
model.nc = num_classes
device = select_device(opt.device, batch_size=batch_size)
model.to(device).train()
criterion = model.loss
optimizer = torch.optim.SGD(model.parameters(), lr=hyp['lr0'], momentum=hyp['momentum'], weight_decay=hyp['weight_decay'])
scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=1, T_mult=2)
start_epoch = 0
best_fitness = 0.0
```
这里使用 Model 类定义了 YOLOv7 模型,并将其放到指定设备上进行训练。使用交叉熵损失函数作为模型的损失函数,使用 SGD 优化器进行训练,并使用余弦退火学习率调整策略。定义了起始轮数、最佳精度等变量。
5. 开始训练
```python
for epoch in range(start_epoch, opt.epochs):
model.train()
mloss = torch.zeros(4).to(device) # mean losses
for i, (imgs, targets, paths, _) in enumerate(train_dataloader):
ni = i + len(train_dataloader) * epoch # number integrated batches (since train start)
imgs = imgs.to(device)
targets = targets.to(device)
loss, _, _ = model(imgs, targets)
loss.backward()
optimizer.step()
optimizer.zero_grad()
mloss = (mloss * i + loss.detach().cpu()) / (i + 1) # update mean losses
# Print batch results
if ni % 20 == 0:
print(f'Epoch {epoch}/{opt.epochs - 1}, Batch {i}/{len(train_dataloader) - 1}, lr={optimizer.param_groups[0]["lr"]:.6f}, loss={mloss[0]:.4f}')
# Update scheduler
scheduler.step()
# Update Best fitness
with torch.no_grad():
fitness = model_fitness(model)
if fitness > best_fitness:
best_fitness = fitness
# Save checkpoint
if (not opt.nosave) or (epoch == opt.epochs - 1):
ckpt = {
'epoch': epoch,
'best_fitness': best_fitness,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict()
}
torch.save(ckpt, f'checkpoints/yolov7_epoch{epoch}.pt')
# Test
if not opt.notest:
t = time_synchronized()
model.eval()
for j, (imgs, targets, paths, shapes) in enumerate(test_dataloader):
if j == 0:
pred = model(imgs.to(device))
pred = non_max_suppression(pred, conf_thres=0.001, iou_thres=0.6)
else:
break
t1 = time_synchronized()
if isinstance(pred, int) or isinstance(pred, tuple):
print(f'Epoch {epoch}/{opt.epochs - 1}, test_loss={mloss[0]:.4f}, test_mAP={0.0}')
else:
pred = pred[0].cpu()
iou_thres = 0.5
niou = [iou_thres] * num_classes
ap, p, r = ap_per_class(pred, targets, shapes, iou_thres=niou)
mp, mr, map50, f1, _, _ = stats(ap, p, r, gt=targets)
print(f'Epoch {epoch}/{opt.epochs - 1}, test_loss={mloss[0]:.4f}, test_mAP={map50:.2f} ({mr*100:.1f}/{mp*100:.1f})')
# Plot images
if epoch == 0 and j == 0:
for i, det in enumerate(pred): # detections per image
img = cv2.imread(paths[i]) # BGR
img = plot_results(img, det, class_names=names)
cv2.imwrite(f'runs/test{i}.jpg', img)
if i == 3:
break
```
这里进行了多个 epoch 的训练。在每个 epoch 中,对于每个批量的数据,先将数据移动到指定设备上,然后计算模型的损失函数,并进行反向传播和梯度下降。在每个 epoch 结束时,更新学习率调整策略和最佳精度,保存当前的检查点。如果 opt.notest 为 False,则进行测试,并输出测试结果。最后,如果是第一个 epoch,则绘制部分图像用于可视化。
阅读全文