從圖像超分辨率快速入門pytorch

時間 2019-11-10

標籤圖像分辨率快速入門 pytorch 简体版

原文原文鏈接

　　前言git

　　最近又開始把pytorch拾起來，學習了github上一些項目以後，發現每一個人都會用不一樣的方式來寫深度學習的訓練代碼，而這些代碼對於初學者來講是難以閱讀的，由於關鍵和非關鍵代碼糅雜在一塊兒，讓那些須要快速將代碼跑起來的初學者摸不着頭腦。github

　　因此，本文打算從最基本的出發，只寫關鍵代碼，將完成一次深度學習訓練須要哪些要素展示給各位初學者，以便大家可以快速上手。等到可以將本身的想法用最簡潔的方式寫出來並運行起來以後，再對本身的代碼進行重構、擴展。我認爲這種學習方式是較好的按部就班的學習方式。網絡

　　本文選擇超分辨率做爲入門案例，一是由於經過結合案例可以對訓練中涉及到的東西有較好的體會，二是超分辨率是較爲簡單的任務，咱們本次教程的目的是教會你們如何使用pytorch，因此不該該將難度設置在任務自己上。下面開始正文。。。多線程

　　正文框架

　　單一圖像超分辨率(SISR)dom

　　簡單介紹一下圖像超分辨率這一任務：超分辨率的任務就是將一張圖像的尺寸放大而且要求失真越小越好，舉例來講，咱們須要將一張256*500的圖像放大2倍，那麼放大後的圖像尺寸就應該是512*1000。用深度學習的方法，咱們一般會先將圖像縮小成原來的1/2，而後以原始圖像做爲標籤，進行訓練。訓練的目標是讓縮小後的圖像放大2倍後與原圖越近越好。因此一般會用L1或者L2做爲損失函數。ide

　　訓練4要素函數

　　一次訓練要想完成，須要的要素我總結爲4點：學習

　　網絡模型優化

　　數據

　　損失函數

　　優化器

　　這4個對象都是一次訓練必不可少的，一般狀況下，須要咱們自定義的是前兩個：網絡模型和數據，然後面兩個較爲統一，並且pytorch也提供了很是全面的實現供咱們使用，它們分別在torch.nn包和torch.optim包下面，使用的時候能夠到pytorch官網進行查看，後面咱們用到的時候還會再次說明。

　　網絡模型

　　在網絡模型和數據兩個當中，網絡模型是比較簡單的，數據加載稍微麻煩些。咱們先來看網絡模型的定義。自定義的網絡模型都必須繼承torch.nn.Module這個類，裏面有兩個方法須要重寫：初始化方法__init__(self)和forward(self, *input)方法。在初始化方法中通常要寫咱們須要哪些層(卷積層、全鏈接層等)，而在forward方法中咱們須要寫這些層的鏈接方式。舉一個通俗的例子，搭積木須要一個個的積木塊，這些積木塊放在__init__方法中，而規定將這些積木塊如何鏈接起來則是靠forward方法中的內容。

　　import torch.nn as nn

　　import torch.nn.functional as F

　　class VDSR(nn.Module):

　　def __init__(self):

　　super(VDSR, self).__init__()

　　self.conv1 = nn.Conv2d(1, 64, kernel_size=3, stride=1, padding=1, bias=True)

　　self.conv2 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=True)

　　self.conv3 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=True)

　　self.conv4 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=True)

　　self.conv5 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=True)

　　self.conv6 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=True)

　　self.conv7 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=True)

　　self.conv8 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=True)

　　self.conv9 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=True)

　　self.conv10 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=True)

　　self.conv11 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=True)

　　self.conv12 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=True)

　　self.conv13 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=True)

　　self.conv14 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=True)

　　self.conv15 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=True)

　　self.conv16 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=True)

　　self.conv17 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=True)

　　self.conv18 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=True)

　　self.conv19 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=True)

　　self.conv20 = nn.Conv2d(64, 1, kernel_size=3, stride=1, padding=1, bias=True)

　　def forward(self, x):

　　ori = x

　　x = F.relu(self.conv1(x))

　　x = F.relu(self.conv2(x))

　　x = F.relu(self.conv3(x))

　　x = F.relu(self.conv4(x))

　　x = F.relu(self.conv5(x))

　　x = F.relu(self.conv6(x))

　　x = F.relu(self.conv7(x))

　　x = F.relu(self.conv8(x))

　　x = F.relu(self.conv9(x))

　　x = F.relu(self.conv10(x))

　　x = F.relu(self.conv11(x))

　　x = F.relu(self.conv12(x))

　　x = F.relu(self.conv13(x))

　　x = F.relu(self.conv14(x))

　　x = F.relu(self.conv15(x))

　　x = F.relu(self.conv16(x))

　　x = F.relu(self.conv17(x))

　　x = F.relu(self.conv18(x))

　　x = F.relu(self.conv19(x))

　　x = self.conv20(x)

　　return x + ori

　　上面代碼中展現的是咱們要用到的模型VDSR，這個模型很簡單，就是連續的20層卷積，外加一個跳線鏈接。結構圖以下：

　　在寫網絡模型時，用到的各個層都在torch.nn這個包中，在寫自定義的網絡結構時能夠自行到pytorch官網的文檔中進行查看。

　　數據

　　定義了網絡模型以後，咱們再來看「數據」。「數據」主要涉及到Dataset和DataLoader兩個概念。

　　Dataset是數據加載的基礎，咱們通常在加載本身的數據集時都須要自定義一個Dataset，自定義的Dataset都須要繼承torch.utils.data.Dataset這個類，當實現了__getitem__()和__len__()這兩個方法後，咱們就自定義了一個Map-style datasets，Dataset是一個可迭代對象，經過下標訪問的方式就可以調用__getitem__()方法來實現數據加載。

　　這裏面最關鍵的就算是__getitem__()如何來寫了，咱們須要讓__getitem__()的返回值是一對，包括圖像和它的label，這裏咱們的任務是超分辨率，那麼圖像和label分別是通過下采樣的圖像和與其對應的原始圖像。因此咱們Dataset的__getitem__()方法返回值就應該是兩個3D Tensor，分別表示兩種圖像。

　　這裏須要重點說明一下__getitem__()方法的返回值爲何應該是3D Tensor。根據pytorch官網的說法，二維卷積層只接受4D Tensor，它的每一維表示的內容分別是nSamples x nChannels x Height x Width，咱們最後須要用批量的方式將數據送到網絡中，因此__getitem__()方法的返回值就應該是後面三維的內容，即使是咱們的通道數爲1，也必須有這一維的存在，不然就會報錯。後面代碼中用到的unsqueeze(0)方法的做用就是如此。前面是說了爲何應該是3D的，爲何應該是Tensor呢?Tensor是跟NumPy中ndarray相似的東西，只是它可以被用於GPU中來加速計算。

　　下面來看一下咱們的代碼：

　　import os

　　import random

　　import cv2

　　import torch

　　from torch.utils.data import Dataset

　　patch_size = 64

　　def getPatch(y):

　　h, w = y.shape

　　randh = random.randrange(0, h - patch_size + 1)

　　randw = random.randrange(0, w - patch_size + 1)

　　lab = y[randh:randh + patch_size, randw:randw + patch_size]

　　resized = cv2.resize(lab, None, fx=0.5, fy=0.5, interpolation=cv2.INTER_CUBIC)

　　rresized = cv2.resize(resized, None, fx=2, fy=2, interpolation=cv2.INTER_CUBIC)

　　return rresized, lab

　　class MyDateSet(Dataset):

　　def __init__(self, imageFolder):

　　self.imageFolder = imageFolder

　　self.images = os.listdir(imageFolder)

　　def __len__(self):

　　return len(self.images)

　　def __getitem__(self, index):

　　name = self.images[index]

　　name = os.path.join(self.imageFolder, name)

　　imread = cv2.imread(name)

　　# 轉換顏色空間

　　ycrcb = cv2.cvtColor(imread, cv2.COLOR_RGB2YCR_CB)

　　# 提取y通道

　　y = ycrcb[:, :, 0]

　　# 裁剪成小塊

　　img, lab = getPatch(y)

　　# 轉爲3D Tensor鄭州婦科醫院 http://www.sptdfk.com/

　　return torch.from_numpy(img).unsqueeze(0), torch.from_numpy(lab).unsqueeze(0)

　　其中MyDateSet的內容也不長，包括了初始化方法、__getitem__()和__len__()兩個方法。__getitem__()有一個輸入值是下標值，咱們根據下標，利用OpenCV，讀取了圖像，並將其轉換顏色空間，超分訓練的時候咱們只用了其中的y通道。還對圖形進行了裁剪，最後返回了兩個3D Tensor。

　　在寫自定義數據集的時候，咱們最須要關注的點就是__getitem__()方法的返回值是否是符合要求，能不可以被送到網絡中去。至於中間該怎麼操做，其實跟pytorch框架也沒什麼關係，根據須要來作。

　　訓練

　　寫好了Dataset以後，咱們就可以經過下標的方式獲取圖像以及它的label。可是離開始訓練還有兩個要素：損失函數和優化器。前面咱們也說了，這兩部分，pytorch官方提供了大量的實現，多數狀況下不須要咱們本身來自定義，這裏咱們直接使用了提供的torch.nn.MSELoss(size_average=None, reduce=None, reduction='mean')做爲損失函數和torch.optim.Adam(params, lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0, amsgrad=False)做爲優化器。

　　訓練示例代碼：

　　import torch

　　import torch.nn as nn

　　import torch.optim as optim

　　import date

　　import model

　　date_set = date.MyDateSet("Train/")

　　model = model.VDSR()

　　device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

　　model.to(device)

　　mse_loss = nn.MSELoss()

　　adam = optim.Adam(model.parameters())

　　for epoch in range(100):

　　running_loss = 0.0

　　for i in range(len(date_set)):

　　rresized, y = date_set[i]

　　adam.zero_grad()

　　out = model(rresized.unsqueeze(0).to(device, torch.float))

　　loss = mse_loss(out, y.unsqueeze(0).to(device, torch.float))

　　loss.backward()

　　adam.step()

　　running_loss += loss

　　if i % 100 == 99: # print every 100

　　print('[%d, %5d] loss: %.3f' %

　　(epoch + 1, i + 1, running_loss / 100))

　　running_loss = 0.0

　　print('Finished Training')

　　整個訓練代碼很是簡潔，只有短短几行，定義模型、將模型移至GPU、定義損失函數、定義優化器(模型移動至GPU必定要在定義優化器以前，由於移動先後的模型已經不是同一個模型對象)。

　　訓練時，先用zero_grad()來將上一次的梯度清零，而後將數據輸入網絡，求偏差，偏差反向傳播求每一個requires_grad=True的Tensor(也就是網絡權重)的梯度，根據優化規則對網絡權重值進行更新，在一次次的更新迭代中，網絡朝着loss下降的方向變化着。

　　值的注意的是，圖像數據也須要移動至GPU，而且須要將其類型轉換爲與網絡模型的權重相同的torch.float

　　DataLoader

　　到前面爲止，其實已經可以實現訓練的過程了，可是，一般狀況下，咱們都須要：

　　將數據打包成一個批量送入網絡

　　每次隨機將數據打亂送入網絡

　　用多線程的方式加載數據(這樣可以提高數據加載速度)

　　這些事情不須要咱們本身實現，有torch.utils.data.DataLoader來幫咱們實現。完整聲明以下：

　　torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, sampler=None, batch_sampler=None, num_workers=0, collate_fn=None, pin_memory=False, drop_last=False, timeout=0, worker_init_fn=None, multiprocessing_context=None)

　　其中的sampler、batch_sampler、collate_fn都是能夠有自定義實現的。咱們簡單的使用默認的實現來構造DataLoader。使用了DataLoader以後的訓練代碼稍微有些不一樣，其中也添加了保存模型的代碼(只保存參數的方式)：

　　import torch

　　import torch.nn as nn

　　import torch.optim as optim

　　from torch.utils.data import DataLoader

　　import date

　　import model

　　date_set = date.MyDateSet("Train/")

　　dataloader = DataLoader(date_set, batch_size=128,

　　shuffle=True, drop_last=True)

　　model = model.VDSR()

　　device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")