Pytorch学习笔记——Dataset

本笔记是根据PyTorch深度学习快速入门教程（绝对通俗易懂！）【小土堆】记录的。

Dataset

1、功能介绍

提供一种方式，获取需要的数据和对应的标签。主要实现两个功能：获取每一个数据及其对应label、统计数据集中的数据数量

2、利用help(Dataset)查看帮助文档

Help on class Dataset in module torch.utils.data.dataset:
class Dataset(typing.Generic)
 |  An abstract class representing a :class:`Dataset`.
 |  
 |  All datasets that represent a map from keys to data samples should subclass
 |  it. All subclasses should overwrite :meth:`__getitem__`, supporting fetching a
 |  data sample for a given key. Subclasses could also optionally overwrite
 |  :meth:`__len__`, which is expected to return the size of the dataset by many
 |  :class:`~torch.utils.data.Sampler` implementations and the default options
 |  of :class:`~torch.utils.data.DataLoader`.
 |  
 |  .. note::
 |    :class:`~torch.utils.data.DataLoader` by default constructs a index
 |    sampler that yields integral indices.  To make it work with a map-style
 |    dataset with non-integral indices/keys, a custom sampler must be provided.
 |  
 |  Method resolution order:
 |      Dataset
 |      typing.Generic
 |      builtins.object
 |  
 |  Methods defined here:
 |  
 |  __add__(self, other: 'Dataset[T_co]') -> 'ConcatDataset[T_co]'
 |  
 |  __getitem__(self, index) -> +T_co
 |  
 |  ----------------------------------------------------------------------
 |  Data descriptors defined here:
 |  
 |  __dict__
 |      dictionary for instance variables (if defined)
 |  
 |  __weakref__
 |      list of weak references to the object (if defined)
 |  
 |  ----------------------------------------------------------------------
 |  Data and other attributes defined here:
 |  
 |  __orig_bases__ = (typing.Generic[+T_co],)
 |  
 |  __parameters__ = (+T_co,)
 |  
 |  ----------------------------------------------------------------------
 |  Class methods inherited from typing.Generic:
 |  
 |  __class_getitem__(params) from builtins.type
 |  
 |  __init_subclass__(*args, **kwargs) from builtins.type
 |      This method is called when a class is subclassed.
 |      
 |      The default implementation does nothing. It may be
 |      overridden to extend subclasses.

这是一个抽象类，所有的数据集想要在数据与标签之间建立映射都需要继承这个类。所有的子类都需要重写__getitem__方法，该方法根据索引值获取每一个数据并且获取其对应的Label。子类也可以重写__len__方法，返回数据集的size大小。

3、实战

导入Dataset类，创建MyData类，继承自Dataset类，使用PIL读取数据

from torch.utils.data import Dataset
from PIL import Image
import os

class MyData(Dataset):
    # 初始化
    def __init__(self, root_dir, label_dir):
        # root_dir数据集根目录文件夹
        self.root_dir = root_dir
        # label_dir为标签
        self.label_dir = label_dir
        # 使用os.path.join()函数，拼接路径
        self.path = os.path.join(self.root_dir, self.label_dir)
        # 将路径下的文件存成数组的形式。数组的元素对应每个图片的名字(str字符串类型)。
        self.img_path = os.listdir(self.path)

    # 对于指定的idx索引，获取数据并返回。
    def __getitem__(self, idx):
        img_name = self.img_path[idx]
        # 拼接数据集目录
        img_item_path = os.path.join(self.root_dir, self.label_dir, img_name)
        # 利用PIL中的Image中的open()函数打开图片
        img = Image.open(img_item_path)
        label = self.label_dir
        return img, label

    # 返回图片文件名列表的长度
    def __len__(self):
        return len(self.img_path)

# 数据集根目录
root_dir = "dataset/train"
# 数据集标签
ants_label_dir = "ants"
bees_label_dir = "bees"
ants_dataset = MyData(root_dir, ants_label_dir)
bees_dataset = MyData(root_dir, bees_label_dir)
# 数据集拼接，当数据集不够的时候可以这么做
train_dataset = ants_dataset + bees_dataset

对每张图像赋予标签写入txt中

import os

root_dir = 'E:/PytorchStudy/dataset/train'
target_dir = 'ants_image'
img_path = os.listdir(os.path.join(root_dir, target_dir))
label = target_dir.split('_')[0]
out_dir = 'ants_label'
for i in img_path:
    file_name = i.split('.jpg')[0]
    with open(os.path.join(root_dir, out_dir,"{}.txt".format(file_name)),'w') as f:
        f.write(label)