【PyTorch图像语义分割】1. PyTorch数据准备与预处理

最新推荐文章于 2024-10-15 21:17:37 发布

通往盛夏

最新推荐文章于 2024-10-15 21:17:37 发布

阅读量8.6k

点赞数 18

分类专栏：数据集 PyTorch 数据加载与预处理文章标签： PyTorch 数据加载与预处理

本文链接：https://blog.csdn.net/tuiqdymy/article/details/84779716

版权

本文档详细介绍了如何在PyTorch中为图像语义分割任务准备数据，包括数据集的结构、数据加载类的创建以及如何进行预处理。通过自定义数据子类并实现数据集大小和索引方法，可以实现图像的批量加载、打乱顺序以及应用变换。此外，代码示例展示了如何处理UAV图像，并提供了简单的测试函数用于检查数据加载的正确性。

摘要由CSDN通过智能技术生成

.py源文件的结构

数据准备与预处理： dataset.py
模型：model.py
训练规则：train.py
测试(benchmark + predict)：test.py

dataset.py的功能

统一将图像返回成torch能处理的[original_iamges.tensor，label.tensor]

用来加载数据的类`torch.utils.data.DataLoader()`

torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, sampler=None, num_workers=0, collate_fn=<function default_collate>, pin_memory=False, drop_last=False)

重点关注四个参数：
batch_size: 批处理数目
shuffle: 是否每个epoch都打乱
workers: 载入数据的线程数
dataset: 是经过变换的自己的数据集（即：一个继承了torch.utils.data.Dataset类的子类的实例），[original_iamges.tensor，label.tensor]之类的，定义的“dataset.py”就是产生这个dataset的。然后在train.py中调用。

导入自己的dataset

class UAVDataSet(torch.ultis.data.Dataset)

继承了torch.utils.data.Dataset这个（抽象）类，我们看看这个类在中文文档中介绍：
所有其他数据集都应该进行子类化。所有子类应该重载__len__和__getitem__，前者提供了数据集的大小，后者支持整数索引，范围从0到len(self)。当然还有个初始化__init__()。
类 = 属性+方法（变量 + 函数），__init__()就是定义自己的属性。

数据子类的基础框架

如上述，必须要重载的是__getitem__()和__len__()。

__len__()：len(dataset)返回数据集的大小。
__getitem__()：实现数据集的下标索引，使用dataset[i]来得到第i个样本（图像和标记）。

import torch.utils.data as data
import torch
from torchvision import transforms
 
class MyTrainData(torch.utils.data.Dataset) #子类化
  def __init__(self, root, transform=None, train=True): #第一步初始化各个变量
 
    self.root = root   
    self.train = train
 
  def