【pytorch学习】 图片数据集的导入和预处理

pytorch提供了一个很方便的类torchvision.datasets.ImageFolder用于读取文件夹中的图片数据并进行预处理。
通过ImageFolder和Dataloader，我们可以快速简单的实现一个数据生成器，每次返回一个batch size的数据和其对应的label。

主要步骤如下：

1. 设置图片预处理参数
2. 用ImageFolder读取数据
3. 用dataloader将数据集分为训练集和测试集，并实现数据生成器

1. 设置图片预处理参数

# 数据预处理，包括数据resize，数据增强（图片翻转等操作），归一化，转成tensor
# 可以直接用imagenet的均值和方差进行normalization
from torchvision import transforms

transform = transforms.Compose(
    [transforms.RandomResizedCrop(input_size),
     transforms.RandomHorizontalFlip(),
     transforms.ToTensor(),
     transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))])
transform_test = transforms.Compose(
    [transforms.Resize(input_size),
     transforms.ToTensor(),
     transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))])

2. 用ImageFolder读取数据

ImageFolder要求数据集有特定的目录结构，举个例子，数据集的地址为‘./animals/dataset’，文件夹的具体目录为：
dataset/dog/xxx.png
dataset/cat/xxx.png

#load image data
data_path = "./animals/dataset"
data_all = torchvision.datasets.ImageFolder(root=data_path, transform=transform)

3. 用dataloader实现数据生成器用于训练

用dataloader将数据集分为训练集和测试集，并实现数据生成器。

dataloader 每次返回一个batch的数据和数据对应的label，使用方法如下：

import torch.utils.data as data
from sklearn.model_selection import train_test_split
from torch.utils.data import SubsetRandomSampler

def dataset_sampler(dataset, val_percentage=0.1):
    """
    split dataset into train set and val set
    :param dataset:
    :param val_percentage: validation percentage
    :return: split sampler
    """
    sample_num = len(dataset)
    file_idx = list(range(sample_num))
    train_idx, val_idx = train_test_split(file_idx, test_size=val_percentage, random_state=42)
    train_sampler = SubsetRandomSampler(train_idx)
    val_sampler = SubsetRandomSampler(val_idx)
    return train_sampler, val_sampler

#将数据集分为训练集和测试集
train_sampler, val_sampler = dataset_sampler(data_all, val_percentage)

#dataloader定义
trainloader = data.DataLoader(data_all, batch_size=batch_size, num_workers=0, sampler=train_sampler)
testloader = data.DataLoader(data_all, batch_size=batch_size, num_workers=0, sampler=val_sampler)

#dataloader用法
for i, data in enumerate(trainloader):
    inputs = data[0].to(device)
    labels = data[1].to(device)

---

最后给一个用resnet进行图像分类的示例：

model = models.resnet34(pretrained=True)
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, 2)

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = model.to(device)
criterion = nn.CrossEntropyLoss()

# Observe that all parameters are being optimized
optimizer_ft = optim.SGD(model.parameters(), lr=learning_rate, momentum=0.9, weight_decay=5e-4)

# Decay LR by a factor of 0.1 every 7 epochs
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=5, gamma=0.1)

train_writer = SummaryWriter(res_dir + '/train')
valid_writer = SummaryWriter(res_dir + '/valid')

def train_model(model, criterion, optimizer, scheduler, num_epochs=10):
    since = time.time()
    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0

    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch, num_epochs - 1))
        print('-' * 10)

        # Each epoch has a training and validation phase
        for phase in ['train', 'val']:
            if phase == 'train':
                scheduler.step()
                model.train()  # Set model to training mode
                dataloader = trainloader
                samples_number = train_num
                # writer = train_writer
            else:
                model.eval()  # Set model to evaluate mode
                dataloader = testloader
                samples_number = test_num
                # writer = valid_writer

            current_loss = 0.0
            current_corrects = 0
            running_loss = 0.0

            # Here's where the training happens
            print('Iterating through data...')

            for i, data in enumerate(dataloader):
                inputs = data[0].to(device)
                labels = data[1].to(device)

                # We need to zero the gradients, don't forget it
                optimizer.zero_grad()

                # Time to carry out the forward training poss
                # We only need to log the loss stats if we are in training phase
                with torch.set_grad_enabled(phase == 'train'):
                    outputs = model(inputs)
                    _, preds = torch.max(outputs, 1)
                    loss = criterion(outputs, labels)

                    # backward + optimize only if in training phase
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                # every 100 batch print loss info
                running_loss += loss.item()
                if i % 100 == 99:
                    print('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, running_loss / 100))
                    running_loss = 0.0

                # We want variables to hold the loss statistics
                current_loss += loss.item() * inputs.size(0)
                current_corrects += (preds == labels).sum().item()

            epoch_loss = current_loss / samples_number
            epoch_acc = current_corrects / samples_number

            print('{} Loss: {:.4f} Acc: {:.4f}'.format(
                phase, epoch_loss, epoch_acc))
            # ...log the running loss
            if phase == 'train':
                train_writer.add_scalar("loss",
                                        epoch_loss,
                                        epoch)
            else:
                valid_writer.add_scalar("loss", epoch_loss, epoch)

        # Make a copy of the model if the accuracy on the validation set has improved
        if phase == 'val' and epoch_acc > best_acc:
            best_acc = epoch_acc
            best_model_wts = copy.deepcopy(model.state_dict())

    print()

    train_writer.close()
    valid_writer.close()
    time_since = time.time() - since
    print('Training complete in {:.0f}m {:.0f}s'.format(
        time_since // 60, time_since % 60))
    print('Best val Acc: {:4f}'.format(best_acc))

    # Now we'll load in the best model weights and return it
    model.load_state_dict(best_model_wts)
    PATH = './CIFAR_net.pth'
    torch.save(best_model_wts, PATH)
    return model


if __name__ == '__main__':
    base_model = train_model(model, criterion, optimizer_ft, exp_lr_scheduler, num_epochs=10)

reference：
https://pytorch.org/tutorials/beginner/data_loading_tutorial.html
https://pytorch.org/tutorials/beginner/finetuning_torchvision_models_tutorial.html#load-data
https://stackabuse.com/image-classification-with-transfer-learning-and-pytorch/