简介
结合官方tutorials和源码以及部分博客写出此文。
pytorch的数据加载和处理相对容易的多，常见的两种形式的导入：
- 一种是整个数据集都在一个文件夹下，内部再另附一个label文件，说明每个文件夹的状态，如这个数据库。这种存放数据的方式可能更适合在非分类问题上得到应用。
- 一种则是更适合使用在分类问题上，即把不同种类的数据分为不同的文件夹存放起来。其形式如下：
root/ants/xxx.png
root/ants/xxy.jpeg
root/ants/xxz.png
.
.
.
root/bees/123.jpg
root/bees/nsdf3.png
root/bees/asd932_.png
本文首先结合官方turorials介绍第一种方法，以了解其数据加载的原理；然后以代码形式简单介绍第二种方法。其中第二种方法和第一种方法的原理相同，其差别在于第二种方法运用了trochvision中提供的已写好的工具ImageFolder，因此实现起来更为简单。
第一种
Dataset class
torch.utils.data.Dataset是一个抽象类，用户想要加载自定义的数据只需要继承这个类，并且覆写其中的两个方法即可：
-
__len__: 覆写这个方法使得len(dataset)可以返回整个数据集的大小 -
__getitem__: 覆写这个方法使得dataset[i]可以返回数据集中第i个样本 - 不覆写这两个方法会直接返回错误，其源码如下：
def __getitem__(self, index):
raise NotImplementedError
def __len__(self):
raise NotImplementedError
这里我随便从网上下载了20张图像，10张小猫，10张小狗。为了省事儿(只是想验证下继承Dataset类是否好用)，我没有给数据集增加标签文件，而是直接把1-10号定义为小猫，11-20号定义为小狗，这样会给__len__和__getitem__减小麻烦，其目录结构如下：
目录结构
建立的自定义类如下：
from torch.utils.data import DataLoader, Dataset
from skimage import io, transform
import matplotlib.pyplot as plt
import os
import torch
from torchvision import transforms
import numpy as np
class AnimalData(Dataset):
def __init__(self, root_dir, transform=None):
self.root_dir = root_dir
self.transform = transform
def __len__(self):
return 20
def __getitem__(self, idx):
filenames = os.listdir(self.root_dir)
filename = filenames[idx]
img = io.imread(os.path.join(self.root_dir, filename))
# print filename[:-5]
if (int(filename[:-5]) > 10):
lable = np.array([0])
else:
lable = np.array([1])
sample = {'image': img, 'lable':lable}
if self.transform:
sample = self.transform(sample)
return sample
Transforms & Compose transforms
可以注意到上一节中AnimalData类中__init__中有个transform参数，这也是这一节中要讲清楚的问题。
从网上随便下载的图片必然大小不一，而cnn的结构却要求输入图像要有固定的大小；numpy中的图像通道定义为H, W, C，而pytorch中的通道定义为C, H, W; pytorch中输入数据需要将numpy array改为tensor类型；输入数据往往需要归一化，等等。
基于以上考虑，我们可以自定义一些Callable的类，然后作为trasform参数传递给上一节定义的dataset类。为了更加方便，torchvision.transforms.Compose提供了Compose类，可以一次性将我们自定义的callable类传递给dataset类，直接得到转换后的数据。
这里我直接copy了教程上的三个类：Rescale, RandomCrop, ToTensor,稍作改动，适应我的数据库。
class Rescale(object):
"""Rescale the image in a sample to a given size.
Args:
output_size (tuple or int): Desired output size. If tuple, output is
matched to output_size. If int, smaller of image edges is matched
to output_size keeping aspect ratio the same.
"""
def __init__(self, output_size):
assert isinstance(output_size, (int, tuple))
self.output_size = output_size
def __call__(self, sample):
image, lable = sample['image'], sample['lable']
h, w = image.shape[:2]
if isinstance(self.output_size, int):
if h > w:
new_h, new_w = self.output_size * h / w, self.output_size
else:
new_h, new_w = self.output_size, self.output_size * w / h
else:
new_h, new_w = self.output_size
new_h, new_w = int(new_h), int(new_w)
img = transform.resize(image, (new_h, new_w))
# h and w are swapped for lable because for images,
# x and y axes are axis 1 and 0 respectively
# lable = lable * [new_w / w, new_h / h]
return {'image': img, 'lable': lable}
class RandomCrop(object):
"""Crop randomly the image in a sample.
Args:
output_size (tuple or int): Desired output size. If int, square crop
is made.
"""
def __init__(self, output_size):
assert isinstance(output_size, (int, tuple))
if isinstance(output_size, int):
self.output_size = (output_size, output_size)
else:
assert len(output_size) == 2
self.output_size = output_size
def __call__(self, sample):
image, lable = sample['image'], sample['lable']
h, w = image.shape[:2]
new_h, new_w = self.output_size
top = np.random.randint(0, h - new_h)
left = np.random.randint(0, w - new_w)
image = image[top: top + new_h,
left: left + new_w]
# lable = lable - [left, top]
return {'image': image, 'lable': lable}
class ToTensor(object):
"""Convert ndarrays in sample to Tensors."""
def __call__(self, sample):
image, lable = sample['image'], sample['lable']
# print lable
# swap color axis because
# numpy image: H x W x C
# torch image: C X H X W
image = image.transpose((2, 0, 1))
return {'image': torch.from_numpy(image),
'lable': torch.from_numpy(lable)}
定义好callable类之后，通过torchvision.transforms.Compose将上述三个类结合在一起，传递给AnimalData类中的transform参数即可。
trsm = transforms.Compose([Rescale(256),
RandomCrop(224),
ToTensor()])
data = AnimalData('./all', transform=trsm)
Iterating through the dataset
上一节中得到data实例之后可以通过for循环来一个一个读取数据，现在这是效率低下的。torch.utils.data.DadaLoader类解决了上述问题。其主要有如下特点：
- Batching the data
- Shuffling the data
- Load the data in parallel using
multiprocessingworkers.
实现起来也很简单：
dataloader = DataLoader(data, batch_size=4, shuffle=True, num_workers=4)
for i_batch, bach_data in enumerate(dataloader):
print i_batch
print bach_data['image'].size()
print bach_data['lable']
第二种
torchvision
pytorch几乎将上述所有工作都封装起来供我们使用，其中一个工具就是torchvision.datasets.ImageFolder,用于加载用户自定义的数据，要求我们的数据要有如下结构：
root/ants/xxx.png
root/ants/xxy.jpeg
root/ants/xxz.png
.
.
.
root/bees/123.jpg
root/bees/nsdf3.png
root/bees/asd932_.png
torchvision.transforms中也封装了各种各样的数据处理的工具，如Resize, ToTensor等等功能供我们使用。
修改我下载的数据库结构如下：
method2_tree
加载数据代码如下：
from torchvision import transforms, utils
from torchvision import datasets
import torch
import matplotlib.pyplot as plt
train_data = datasets.ImageFolder('./data1', transform=transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor()
]))
train_loader = torch.utils.data.DataLoader(train_data,
batch_size=4,
shuffle=True,
)
print len(train_loader)
for i_batch, img in enumerate(train_loader):
if i_batch == 0:
print(img[1])
fig = plt.figure()
grid = utils.make_grid(img[0])
plt.imshow(grid.numpy().transpose((1, 2, 0)))
plt.show()
break
结果图：
make_grid
附录
最后欣赏一段torchvision源码
# vision/torchvision/datasets/folder.py
import torch.utils.data as data
from PIL import Image
import os
import os.path
IMG_EXTENSIONS = ['.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm']
def is_image_file(filename):
"""Checks if a file is an image.
Args:
filename (string): path to a file
Returns:
bool: True if the filename ends with a known image extension
"""
filename_lower = filename.lower()
return any(filename_lower.endswith(ext) for ext in IMG_EXTENSIONS)
def find_classes(dir):
classes = [d for d in os.listdir(dir) if os.path.isdir(os.path.join(dir, d))]
classes.sort()
class_to_idx = {classes[i]: i for i in range(len(classes))}
return classes, class_to_idx
def make_dataset(dir, class_to_idx):
images = []
dir = os.path.expanduser(dir)
for target in sorted(os.listdir(dir)):
d = os.path.join(dir, target)
if not os.path.isdir(d):
continue
for root, _, fnames in sorted(os.walk(d)):
for fname in sorted(fnames):
if is_image_file(fname):
path = os.path.join(root, fname)
item = (path, class_to_idx[target])
images.append(item)
return images
def pil_loader(path):
# open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835)
with open(path, 'rb') as f:
img = Image.open(f)
return img.convert('RGB')
def accimage_loader(path):
import accimage
try:
return accimage.Image(path)
except IOError:
# Potentially a decoding problem, fall back to PIL.Image
return pil_loader(path)
def default_loader(path):
from torchvision import get_image_backend
if get_image_backend() == 'accimage':
return accimage_loader(path)
else:
return pil_loader(path)
class ImageFolder(data.Dataset):
"""A generic data loader where the images are arranged in this way: ::
root/dog/xxx.png
root/dog/xxy.png
root/dog/xxz.png
root/cat/123.png
root/cat/nsdf3.png
root/cat/asd932_.png
Args:
root (string): Root directory path.
transform (callable, optional): A function/transform that takes in an PIL image
and returns a transformed version. E.g, ``transforms.RandomCrop``
target_transform (callable, optional): A function/transform that takes in the
target and transforms it.
loader (callable, optional): A function to load an image given its path.
Attributes:
classes (list): List of the class names.
class_to_idx (dict): Dict with items (class_name, class_index).
imgs (list): List of (image path, class_index) tuples
"""
def __init__(self, root, transform=None, target_transform=None,
loader=default_loader):
classes, class_to_idx = find_classes(root)
imgs = make_dataset(root, class_to_idx)
if len(imgs) == 0:
raise(RuntimeError("Found 0 images in subfolders of: " + root + "\n"
"Supported image extensions are: " + ",".join(IMG_EXTENSIONS)))
self.root = root
self.imgs = imgs
self.classes = classes
self.class_to_idx = class_to_idx
self.transform = transform
self.target_transform = target_transform
self.loader = loader
def __getitem__(self, index):
"""
Args:
index (int): Index
Returns:
tuple: (image, target) where target is class_index of the target class.
"""
path, target = self.imgs[index]
img = self.loader(path)
if self.transform is not None:
img = self.transform(img)
if self.target_transform is not None:
target = self.target_transform(target)
return img, target
def __len__(self):
return len(self.imgs)
def __repr__(self):
fmt_str = 'Dataset ' + self.__class__.__name__ + '\n'
fmt_str += ' Number of datapoints: {}\n'.format(self.__len__())
fmt_str += ' Root Location: {}\n'.format(self.root)
tmp = ' Transforms (if any): '
fmt_str += '{0}{1}\n'.format(tmp, self.transform.__repr__().replace('\n', '\n' + ' ' * len(tmp)))
tmp = ' Target Transforms (if any): '
fmt_str += '{0}{1}'.format(tmp, self.target_transform.__repr__().replace('\n', '\n' + ' ' * len(tmp)))
return fmt_str
参考
[1]. Data Loading and Processing Tutorial
[2]. github: pytorch/torch/utils/data/dataset.py
[3]. github: vision/torchvision/datasets/folder.py
[4]. csdn