轉載：

pytorch載入資料與預處理資料 - pytorch中文網
原文出處： https://ptorch.com/news/140.html

解決任何機器學習問題需要付出很多努力來準備資料。PyTorch提供了許多工具可以使資料載入變得輕鬆而有希望，從而使您的程式碼更具可讀性。在本教程中，我們將看到如何對不一般的資料進行載入和預處理/資料增強。

要執行本教程，請確保已安裝以下的安裝包：

• scikit-image：用於影象io和transforms
• pandas：為了更簡單的解析csv

from __future__ import print_function, division
import os
import torch
import pandas as pd
from skimage import io, transform
import numpy as np
import matplotlib.pyplot as plt
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, utils

# Ignore warnings
import warnings
warnings.filterwarnings("ignore")

plt.ion()   # interactive mode
 

我們要處理的資料集是面部姿勢。意思是這樣標註一張臉：

每張臉的影象總共有68個不同的標記點。

注意 你可以從這裡下載資料集，這份資料是從ImageNet中選取一些標記為face的圖片，使用dlib’s pose estimation方法生成的。

資料集裡面附帶一個註釋如下所示的csv檔案：

image_name,part_0_x,part_0_y,part_1_x,part_1_y,part_2_x, ... ,part_67_x,part_67_y
0805personali01.jpg,27,83,27,98, ... 84,134
1084239450_e76e00b7e7.jpg,70,236,71,257, ... ,128,312
 

讓我們快速讀取CSV檔案，並且以(N, 2)的陣列形式獲得標記點，其中N標記點的個數。

landmarks_frame = pd.read_csv('faces/face_landmarks.csv')

n = 65
img_name = landmarks_frame.iloc[n, 0]
landmarks = landmarks_frame.iloc[n, 1:].as_matrix()
landmarks = landmarks.astype('float').reshape(-1, 2)

print('Image name: {}'.format(img_name))
print('Landmarks shape: {}'.format(landmarks.shape))
print('First 4 Landmarks: {}'.format(landmarks[:4]))
 

輸出：

Image name: person-7.jpg
Landmarks shape: (68, 2)
First 4 Landmarks: [[ 32.  65.]
 [ 33.  76.]
 [ 34.  86.]
 [ 34.  97.]]

讓我們寫一個簡單的幫助函式來顯示影象和標記點，並用它來顯示一個樣本。

def show_landmarks(image, landmarks):
    """Show image with landmarks"""
    plt.imshow(image)
    plt.scatter(landmarks[:, 0], landmarks[:, 1], s=10, marker='.', c='r')
    plt.pause(0.001)  # pause a bit so that plots are updated

plt.figure()
show_landmarks(io.imread(os.path.join('faces/', img_name)),
               landmarks)
plt.show()

資料集類

torch.utils.data.Dataset是表示資料集的抽象類。您自定義的資料集應該繼承Dataset並重寫以下方法：

• __len__ 使用len(dataset)將返回資料集的大小。
• __getitem__ 支援索引，dataset[i]可以獲取第i個樣本

讓我們為我們的人臉標記點資料集建立一個數據集類。我們將在__init__中讀取csv，而將在__getitem__存放讀取圖片的任務。因為所有的影象不是一次性儲存在記憶體中，而是根據需要進行讀取，這樣可以高效的使用記憶體。

我們的資料集是一個字典{'image': image, 'landmarks': landmarks}。資料集的類有一個可選的引數transform，這樣就可以對資料做特定的預處理操作。在下一節中我們會看到transfrom的作用。

class FaceLandmarksDataset(Dataset):
    """Face Landmarks dataset."""

    def __init__(self, csv_file, root_dir, transform=None):
        """
        Args:
            csv_file (string): Path to the csv file with annotations.
            root_dir (string): Directory with all the images.
            transform (callable, optional): Optional transform to be applied
                on a sample.
        """
        self.landmarks_frame = pd.read_csv(csv_file)
        self.root_dir = root_dir
        self.transform = transform

    def __len__(self):
        return len(self.landmarks_frame)

    def __getitem__(self, idx):
        img_name = os.path.join(self.root_dir,
                                self.landmarks_frame.iloc[idx, 0])
        image = io.imread(img_name)
        landmarks = self.landmarks_frame.iloc[idx, 1:].as_matrix()
        landmarks = landmarks.astype('float').reshape(-1, 2)
        sample = {'image': image, 'landmarks': landmarks}

        if self.transform:
            sample = self.transform(sample)

        return sample

讓我們例項化這個類，並且遍歷所有的資料樣本。我們將列印前4個樣本的形狀，顯示他們的標記點。

face_dataset = FaceLandmarksDataset(csv_file='faces/face_landmarks.csv',
                                    root_dir='faces/')

fig = plt.figure()

for i in range(len(face_dataset)):
    sample = face_dataset[i]

    print(i, sample['image'].shape, sample['landmarks'].shape)

    ax = plt.subplot(1, 4, i + 1)
    plt.tight_layout()
    ax.set_title('Sample #{}'.format(i))
    ax.axis('off')
    show_landmarks(**sample)

    if i == 3:
        plt.show()
        break

輸出：

0 (324, 215, 3) (68, 2)
1 (500, 333, 3) (68, 2)
2 (250, 258, 3) (68, 2)
3 (434, 290, 3) (68, 2)

Transforms

從上面我們可以看到的一個問題是樣本的大小不一樣。大多數神經網路期望固定大小的影象。因此，我們需要編寫一些預處理程式碼。我們建立三個轉換：

• Rescale：縮放影象
• RandomCrop：隨機從影象中裁剪。這是資料增強。
• ToTensor：將numpy影象轉換為torch影象（我們需要交換維度）。

我們把它們寫成可以呼叫的類而不是簡單的函式，這樣每次只需要傳遞需要的引數就可以呼叫transforms函數了。這樣的話，我們就需要完成__call__方法以及__init__，如果需要這個的話。之後我們就可以像這樣使用transforms了：

tsfm = Transform(params)
transformed_sample = tsfm(sample)

仔細觀察下面的函式是如何同時對圖片和標記點做transforms的。

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, landmarks = sample['image'], sample['landmarks']

        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 landmarks because for images,
        # x and y axes are axis 1 and 0 respectively
        landmarks = landmarks * [new_w / w, new_h / h]

        return {'image': img, 'landmarks': landmarks}

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, landmarks = sample['image'], sample['landmarks']

        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]

        landmarks = landmarks - [left, top]

        return {'image': image, 'landmarks': landmarks}

class ToTensor(object):
    """Convert ndarrays in sample to Tensors."""

    def __call__(self, sample):
        image, landmarks = sample['image'], sample['landmarks']

        # 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),
                'landmarks': torch.from_numpy(landmarks)}

Compose transforms

現在我們就將transform運用在一個樣本上。

我們說我們想將圖片的短邊變為256，之後隨機裁剪一個邊長為224的正方形。這樣的話，我們就需要組合Rescale和RandomCrop。torchvision.transforms.Compose是一個簡單的可呼叫的類，它允許我們像下面這樣做。

scale = Rescale(256)
crop = RandomCrop(128)
composed = transforms.Compose([Rescale(256),
                               RandomCrop(224)])

# Apply each of the above transforms on sample.
fig = plt.figure()
sample = face_dataset[65]
for i, tsfrm in enumerate([scale, crop, composed]):
    transformed_sample = tsfrm(sample)

    ax = plt.subplot(1, 3, i + 1)
    plt.tight_layout()
    ax.set_title(type(tsfrm).__name__)
    show_landmarks(**transformed_sample)

plt.show()

遍歷資料集

我們將所有的放在一起來生成composed transforms之後的資料。總之，每次迭代的資料:

• 從檔案中讀取一幅影象
• 在讀取圖片時進行transforms操作
• 由於其中一個transforms是隨機的，所以迭代的資料樣本進行了增強

我們可以像之前一樣使用for i in range迴圈從建立的資料中進行迭代。

transformed_dataset = FaceLandmarksDataset(csv_file='faces/face_landmarks.csv',
                                           root_dir='faces/',
                                           transform=transforms.Compose([
                                               Rescale(256),
                                               RandomCrop(224),
                                               ToTensor()
                                           ]))

for i in range(len(transformed_dataset)):
    sample = transformed_dataset[i]

    print(i, sample['image'].size(), sample['landmarks'].size())

    if i == 3:
        break

輸出：

0 torch.Size([3, 224, 224]) torch.Size([68, 2])
1 torch.Size([3, 224, 224]) torch.Size([68, 2])
2 torch.Size([3, 224, 224]) torch.Size([68, 2])
3 torch.Size([3, 224, 224]) torch.Size([68, 2])

然而，通過使用一個簡單的for迴圈來迭代資料我們可能損失很多資訊。尤其是我們丟失了這些操作：

• 批量處理資料
• 打亂資料順序
• 使用多程序（multiprocessing）並行載入資料

torch.utils.data.DataLoader是一個提供所有上面這些功能的迭代器。下面使用的引數應該清楚。其中一個蠻有趣的引數是collate_fn。你可以使用collate_fn來指定如何讀取一批的額樣本。然而，預設的collate在大部分的情況下都表現得很好。

dataloader = DataLoader(transformed_dataset, batch_size=4,
                        shuffle=True, num_workers=4)

# Helper function to show a batch
def show_landmarks_batch(sample_batched):
    """Show image with landmarks for a batch of samples."""
    images_batch, landmarks_batch = \
            sample_batched['image'], sample_batched['landmarks']
    batch_size = len(images_batch)
    im_size = images_batch.size(2)

    grid = utils.make_grid(images_batch)
    plt.imshow(grid.numpy().transpose((1, 2, 0)))

    for i in range(batch_size):
        plt.scatter(landmarks_batch[i, :, 0].numpy() + i * im_size,
                    landmarks_batch[i, :, 1].numpy(),
                    s=10, marker='.', c='r')

        plt.title('Batch from dataloader')

for i_batch, sample_batched in enumerate(dataloader):
    print(i_batch, sample_batched['image'].size(),
          sample_batched['landmarks'].size())

    # observe 4th batch and stop.
    if i_batch == 3:
        plt.figure()
        show_landmarks_batch(sample_batched)
        plt.axis('off')
        plt.ioff()
        plt.show()
        break

輸出：

0 torch.Size([4, 3, 224, 224]) torch.Size([4, 68, 2])
1 torch.Size([4, 3, 224, 224]) torch.Size([4, 68, 2])
2 torch.Size([4, 3, 224, 224]) torch.Size([4, 68, 2])
3 torch.Size([4, 3, 224, 224]) torch.Size([4, 68, 2])

後記：torchvision

在本教程中，我們已經看到如何編寫和使用資料集，transforms和dataloader。torchvision包提供了一些常用的資料集和transforms。你甚至可能不需要編寫自定義類。torchvision提供的更通用的資料集之一是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

其中，ants, bees等是類別。對PIL.Image進行的通用的transforms操作如RandomHorizontalFlip, Scale也可以隨時使用。你可以像這樣使用這樣函式來寫dataloader：

import torch
from torchvision import transforms, datasets

data_transform = transforms.Compose([
        transforms.RandomSizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                             std=[0.229, 0.224, 0.225])
    ])
hymenoptera_dataset = datasets.ImageFolder(root='hymenoptera_data/train',
                                           transform=data_transform)
dataset_loader = torch.utils.data.DataLoader(hymenoptera_dataset,
                                             batch_size=4, shuffle=True,
                                             num_workers=4)