PyTorch 学习笔记(一)：transforms&imgaug

PyTorch 学习笔记(一)：transforms&imgaug

1.transforms

https://blog.csdn.net/u011995719/article/details/85107009

2. imgaug

https://github.com/aleju/imgaug

https://imgaug.readthedocs.io/en/latest/

https://github.com/mdbloice/Augmentor

https://github.com/albu/albumentations

3. Augment Demo

# --coding:utf-8--
import numpy as np
import matplotlib.pyplot as plt
import os
import torch
import torch.nn as nn
import torchvision
from torchvision import models, transforms, datasets
import time
from tqdm import tqdm
import sys
import matplotlib.pyplot as plt
import PIL
from imgaug import augmenters as iaa
import imgaug as ia

# 检查GPU是否可用
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

def reg2class(y, num):
    cls_y = torch.zeros(y.shape[0], num)
    for i in range(y.shape[0]):
        cls_y[i, y[i]] = 1
    return cls_y

class ImgAugTransform:
    def __init__(self):
        self.aug = iaa.Sequential([
            iaa.Scale((224, 224)),
            iaa.Sometimes(0.25, iaa.GaussianBlur(sigma=(0, 3.0))),
            iaa.Fliplr(0.5),
            iaa.Affine(rotate=(-10, 10), mode='symmetric'),
            iaa.Sometimes(0.25,
                          iaa.OneOf([iaa.Dropout(p=(0, 0.1)),
                                     iaa.CoarseDropout(0.1, size_percent=0.5)])),
            iaa.AddToHueAndSaturation(value=(-10, 10), per_channel=True)
        ])

    def __call__(self, img):
        img = np.array(img)
        return self.aug.augment_image(img)


class ResNet34C3(nn.Module):
    def __init__(self, n_classes):
        super(ResNet34C3, self).__init__()
        resnet = models.resnet34(pretrained=True)
        # 是否返回梯度/固定参数
        for p in resnet.parameters():
            p.requires_grad = False
        self.feature = resnet
        self.nclass = n_classes
        self.relu = nn.ReLU()
        self.fc1 = nn.Linear(512, 256)
        self.bn1 = nn.BatchNorm1d(256)
        self.fc2 = nn.Linear(256, 64)
        self.bn2 = nn.BatchNorm1d(64)
        self.fc3 = nn.Linear(64, self.nclass)

    def forward(self, x):
        # 1000
        x = self.feature.conv1(x)
        x = self.feature.bn1(x)
        x = self.feature.relu(x)
        x = self.feature.maxpool(x)
        x = self.feature.layer1(x)
        x = self.feature.layer2(x)
        x = self.feature.layer3(x)
        x = self.feature.layer4(x)
        x = self.feature.avgpool(x)
        x = x.view(x.size(0), -1)
        x = self.fc1(x)
        X = self.bn1(x)
        x = self.relu(x)
        x = self.fc2(x)
        X = self.bn2(x)
        x = self.relu(x)
        x = self.fc3(x)
        # forward 参数必须初始化，未初始化GPU训练出错，CPU可训练
        # x = nn.Linear(1000, 512, bias=True)(x)
        # x = nn.LeakyReLU(0.1)(x)
        # x = nn.Linear(512, 256, bias=True)(x)
        # x = nn.LeakyReLU(0.1)(x)
        # x = nn.Linear(256, 64, bias=True)(x)
        # x = nn.LeakyReLU(0.1)(x)
        # x = nn.Dropout(p=0.5, inplace=False)(x)
        # x = nn.Linear(512, 3)(x)

        return x


if __name__ == '__main__':
    # 数据集路径
    data_dir = '../../DataSet'
    EPOCHS = 1000
    LR = 0.002
    numClass = 3
    # 数据输入格式
    # res_format = transforms.Compose([
    #     transforms.Resize((224, 224)),
    #     transforms.ColorJitter(hue=.5, saturation=.5),
    #     transforms.RandomHorizontalFlip(),
    #     transforms.RandomRotation(10, resample=PIL.Image.BILINEAR),
    #     transforms.RandomAffine(10, translate=None),
    #     transforms.ToTensor()
    # ])
    # pytorch Channel first
    # dsets = {x: datasets.ImageFolder(os.path.join(data_dir, x),res_format )for x in ['train', 'valid']}
    # imgaug Channel Last
    ia_trans = ImgAugTransform()
    dsets = {x: datasets.ImageFolder(os.path.join(data_dir, x), ia_trans)for x in ['train', 'valid']}
    dset_sizes = {x: len(dsets[x]) for x in ['train', 'valid']}
    print(dset_sizes)
    dset_classes = dsets['train'].classes
    loader_train = torch.utils.data.DataLoader(dsets['train'], batch_size=64, shuffle=True, num_workers=6)
    loader_valid = torch.utils.data.DataLoader(dsets['valid'], batch_size=10, shuffle=True, num_workers=6)
    # Model
    SDG_Net = ResNet34C3(numClass)
    if torch.cuda.is_available():
        SDG_Net = SDG_Net.cuda()
    criterion = torch.nn.CrossEntropyLoss()
    optim = torch.optim.Adam(SDG_Net.parameters(), lr=LR, weight_decay=0.0)
    # 开始训练
    for epoch in tqdm(range(EPOCHS)):
        index = 0
        # 训练过程调整学习率
        if epoch % 100 == 0:
            for param_group in optim.param_groups:
                LR = LR * 0.9
                param_group['lr'] = LR
        bi = 0
        for data in loader_train:
            batch_x, batch_y = data
            batch_x = batch_x.float()
            batch_x = batch_x.transpose(1, -1)
            # clsy = reg2class(batch_y, numClass)
            # plt.imshow(batch_x[5, 0].cpu().numpy(), cmap='gray')
            if torch.cuda.is_available():
                batch_x = batch_x.cuda()
                batch_y = batch_y.cuda()
            prediction = SDG_Net.forward(batch_x)
            train_loss = criterion(prediction, batch_y)
            optim.zero_grad()
            train_loss.backward()
            optim.step()
            train_info = '\r>>epoch:' + str(epoch) + ', batch_num:' + str(bi) + ', train_loss:' + str(
                train_loss.cpu().detach().numpy())
            sys.stdout.write(train_info)
            sys.stdout.flush()
            # print(train_info)
            bi = bi + 1
            # print('train_loss:%4f' % train_loss)
        for data in loader_valid:
            batch_x, batch_y = data
            batch_x = batch_x.float()
            batch_x = batch_x.transpose(1, -1)
            # clsy = reg2class(batch_y, numClass)
            if torch.cuda.is_available():
                batch_x = batch_x.cuda()
                batch_y = batch_y.cuda()
            # batch_x = torch.tensor(batch_x, device=device)
            # batch_y = torch.tensor(batch_y, device=device)
            prediction = SDG_Net.forward(batch_x)
            valid_loss = criterion(prediction, batch_y)
            # optim.zero_grad()
            # loss.backward()
            # optim.step()
        train_info = 'epoch:' + str(epoch) + ', train_loss:' + str(train_loss.cpu().detach().numpy()) + \
                     ', valid_loss:' + str(valid_loss.cpu().detach().numpy())
        print(train_info)
        torch.save(SDG_Net.state_dict(), '../../TrainLog/' + "SDG_Net_%d_train_%4f_valid_%4f.pth" % (
            epoch, train_loss, valid_loss))

    print('End-----------------------')