PyTorch入門代碼——訓練一個圖像分類模型(Level1)

7 人贊了文章

這是一個適合PyTorch入門者看的博客。PyTorch的文檔質量比較高，入門較為容易，這篇博客選取[官方鏈接](http://pytorch.org/tutorials/beginner/transfer_learning_tutorial.html)裡面的例子，介紹如何用PyTorch訓練一個ResNet模型用於圖像分類，代碼邏輯非常清晰，基本上和許多深度學習框架的代碼思路類似，非常適合初學者想上手PyTorch訓練模型（不必每次都跑mnist的demo了）。接下來從個人使用角度加以解釋。解釋的思路是從數據導入開始到模型訓練結束，基本上就是搭積木的方式來寫代碼。

首先是數據導入部分，這裡採用官方寫好的torchvision.datasets.ImageFolder介面實現數據導入。這個介面需要你提供圖像所在的文件夾，就是下面的data_dir=『/data』這句，然後對於一個分類問題，這裡data_dir目錄下一般包括兩個文件夾：train和val，每個文件件下面包含N個子文件夾，N是你的分類類別數，且每個子文件夾里存放的就是這個類別的圖像。這樣torchvision.datasets.ImageFolder就會返回一個列表（比如下面代碼中的image_datasets[train]或者image_datasets[val]），列表中的每個值都是一個tuple，每個tuple包含圖像和標籤信息。

data_dir = /dataimage_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x]) for x in [train, val]}

另外這裡的data_transforms是一個字典，如下。主要是進行一些圖像預處理，比如resize、crop等。實現的時候採用的是torchvision.transforms模塊，比如torchvision.transforms.Compose是用來管理所有transforms操作的，torchvision.transforms.RandomSizedCrop是做crop的。需要注意的是對於torchvision.transforms.RandomSizedCrop和transforms.RandomHorizontalFlip()等，輸入對象都是PIL Image，也就是用python的PIL庫讀進來的圖像內容，而transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])的作用對象需要是一個Tensor，因此在transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])之前有一個 transforms.ToTensor()就是用來生成Tensor的。另外transforms.Scale(256)其實就是resize操作，目前已經被transforms.Resize類取代了。

data_transforms = { train: transforms.Compose([ transforms.RandomSizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]), val: transforms.Compose([ transforms.Scale(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]), }

前面torchvision.datasets.ImageFolder只是返回list，list是不能作為模型輸入的，因此在PyTorch中需要用另一個類來封裝list，那就是：torch.utils.data.DataLoader。torch.utils.data.DataLoader類可以將list類型的輸入數據封裝成Tensor數據格式，以備模型使用。注意，這裡是對圖像和標籤分別封裝成一個Tensor。這裡要提到另一個很重要的類：torch.utils.data.Dataset，這是一個抽象類，在pytorch中所有和數據相關的類都要繼承這個類來實現。比如前面說的torchvision.datasets.ImageFolder類是這樣的，以及這裡的torch.util.data.DataLoader類也是這樣的。所以當你的數據不是按照一個類別一個文件夾這種方式存儲時，你就要自定義一個類來讀取數據，自定義的這個類必須繼承自torch.utils.data.Dataset這個基類，最後同樣用torch.utils.data.DataLoader封裝成Tensor。

dataloders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=4, shuffle=True, num_workers=4) for x in [train, val]}

生成dataloaders後再有一步就可以作為模型的輸入了，那就是將Tensor數據類型封裝成Variable數據類型，來看下面這段代碼。dataloaders是一個字典，dataloders[train]存的就是訓練的數據，這個for循環就是從dataloders[train]中讀取batch_size個數據，batch_size在前面生成dataloaders的時候就設置了。因此這個data裡面包含圖像數據（inputs）這個Tensor和標籤（labels）這個Tensor。然後用torch.autograd.Variable將Tensor封裝成模型真正可以用的Variable數據類型。 為什麼要封裝成Variable呢？在pytorch中，torch.tensor和torch.autograd.Variable是兩種比較重要的數據結構，Variable可以看成是tensor的一種包裝，其不僅包含了tensor的內容，還包含了梯度等信息，因此在神經網路中常常用Variable數據結構。那麼怎麼從一個Variable類型中取出tensor呢？也很簡單，比如下面封裝後的inputs是一個Variable，那麼inputs.data就是對應的tensor。

for data in dataloders[train]: inputs, labels = data if use_gpu: inputs = Variable(inputs.cuda()) labels = Variable(labels.cuda()) else: inputs, labels = Variable(inputs), Variable(labels)

封裝好了數據後，就可以作為模型的輸入了。所以要先導入你的模型。在PyTorch中已經默認為大家準備了一些常用的網路結構，比如分類中的VGG，ResNet，DenseNet等等，可以用torchvision.models模塊來導入。比如用torchvision.models.resnet18(pretrained=True)來導入ResNet18網路，同時指明導入的是已經預訓練過的網路。因為預訓練網路一般是在1000類的ImageNet數據集上進行的，所以要遷移到你自己數據集的2分類，需要替換最後的全連接層為你所需要的輸出。因此下面這三行代碼進行的就是用models模塊導入resnet18網路，然後獲取全連接層的輸入channel個數，用這個channel個數和你要做的分類類別數（這裡是2）替換原來模型中的全連接層。這樣網路結果也準備好。

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

但是只有網路結構和數據還不足以讓代碼運行起來，還需要定義損失函數。在PyTorch中採用torch.nn模塊來定義網路的所有層，比如卷積、降採樣、損失層等等，這裡採用交叉熵函數，因此可以這樣定義：

criterion = nn.CrossEntropyLoss()

然後你還需要定義優化函數，比如最常見的隨機梯度下降，在PyTorch中是通過torch.optim模塊來實現的。另外這裡雖然寫的是SGD，但是因為有momentum，所以是Adam的優化方式。這個類的輸入包括需要優化的參數：model.parameters()，學習率，還有Adam相關的momentum參數。現在很多優化方式的默認定義形式就是這樣的。

optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

然後一般還會定義學習率的變化策略，這裡採用的是torch.optim.lr_scheduler模塊的StepLR類，表示每隔step_size個epoch就將學習率降為原來的gamma倍。

scheduler = lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)

準備工作終於做完了，要開始訓練了。

首先訓練開始的時候需要先更新下學習率，這是因為我們前面制定了學習率的變化策略，所以在每個epoch開始時都要更新下：

scheduler.step()

然後設置模型狀態為訓練狀態：

model.train(True)

然後先將網路中的所有梯度置0：

optimizer.zero_grad()

然後就是網路的前向傳播了：

outputs = model(inputs)

然後將輸出的outputs和原來導入的labels作為loss函數的輸入就可以得到損失了：

loss = criterion(outputs, labels)

輸出的outputs也是torch.autograd.Variable格式，得到輸出後（網路的全連接層的輸出）還希望能到到模型預測該樣本屬於哪個類別的信息，這裡採用torch.max。torch.max()的第一個輸入是tensor格式，所以用outputs.data而不是outputs作為輸入；第二個參數1是代表dim的意思，也就是取每一行的最大值，其實就是我們常見的取概率最大的那個index；第三個參數loss也是torch.autograd.Variable格式。

_, preds = torch.max(outputs.data, 1)

計算得到loss後就要回傳損失。要注意的是這是在訓練的時候才會有的操作，測試時候只有forward過程。

loss.backward()

回傳損失過程中會計算梯度，然後需要根據這些梯度更新參數，optimizer.step()就是用來更新參數的。optimizer.step()後，你就可以從optimizer.param_groups[0][params]裡面看到各個層的梯度和權值信息。

optimizer.step()

這樣一個batch數據的訓練就結束了！當你不斷重複這樣的訓練過程，最終就可以達到你想要的結果了。

另外如果你有gpu可用，那麼包括你的數據和模型都可以在gpu上操作，這在PyTorch中也非常簡單。判斷你是否有gpu可以用可以通過下面這行代碼，如果有，則use_gpu是true。

use_gpu = torch.cuda.is_available()

完整代碼請移步：Github

推薦閱讀：