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-----------------------')