PyTorch常用代碼

簡介

之前提到過PyTorch進行深度訓(xùn)練是一套總的來說較為固定的流程，因此學(xué)習(xí)一些常用的代碼段是很有必要的，本文整理了一些常用的代碼段，當(dāng)然最合適的參考資料依然是官方文檔。

框架配置

包導(dǎo)入與信息查看

這部分主要是PyTorch的版本查看、本機是否可以使用GPU訓(xùn)練、可用的GPU數(shù)目等，具體如下代碼，已經(jīng)備注詳細。

import torch# PyTorch版本 print(torch.__version__) # 是否可用GPU print(torch.cuda.is_available()) # 可用的GPU數(shù)目 print(torch.cuda.device_count()) # 可用的第一個GPU（默認從0編號開始） print(torch.cuda.get_device_name(0) if torch.cuda.device_count() > 0 else "no gpu")

固定隨機數(shù)種子

在硬件設(shè)備不同的情況下，完整的復(fù)現(xiàn)是比較困難的，但是同一個設(shè)備上應(yīng)該盡可能保證可復(fù)現(xiàn)性，主要做法就是在程序開始之前固定隨機數(shù)種子，包括Numpy、Torch的隨機數(shù)種子。

import numpy as np import torchnp.random.seed(0) torch.manual_seed(0) torch.cuda.manual_seed_all(0)

顯卡設(shè)置

單顯卡設(shè)置如下。

import torch device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

指定多張顯卡設(shè)置（下面代碼示例指定0和1號顯卡）。

import os os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'

張量運算

PyTorch總共提供了9種CPU張量類型以及對應(yīng)的GPU張量類型，這部分主要是張量的一些常用操作，不包括基本的張量創(chuàng)建之類的API。

張量信息

對于一個張量，有時代碼調(diào)試時關(guān)注其基本信息如維度、數(shù)據(jù)類型。

import torcht = torch.randn(2, 3, 4, 4) # 張量的數(shù)據(jù)類型 print(t.type()) # 張量的維度信息 print(t.size()) # 張量的維度數(shù)量 print(t.dim())

數(shù)據(jù)類型轉(zhuǎn)換

PyTorch經(jīng)常涉及到不同數(shù)據(jù)類型之間的轉(zhuǎn)換，尤其是CPU數(shù)據(jù)轉(zhuǎn)為GPU數(shù)據(jù)，有時候為了方便運行也需要設(shè)置默認的數(shù)據(jù)類型（在PyTorch中Float類型遠超Double類型的速度）。

import torch# 設(shè)置默認數(shù)據(jù)類型 torch.set_default_tensor_type(torch.FloatTensor)t = torch.randn(1) print(t.type()) # 轉(zhuǎn)為GPU數(shù)據(jù)類型 t = t.cuda() # 轉(zhuǎn)為CPU數(shù)據(jù)類型 t = t.cpu() t = t.float() t = t.long()

Numpy轉(zhuǎn)換

除了字符類型，其他類型的張量都支持轉(zhuǎn)換為Numpy的數(shù)組，也支持Numpy的數(shù)組轉(zhuǎn)為PyTorch的張量。

import torcht = torch.randn(1, 2, 3) # Tensor轉(zhuǎn)ndarray t = t.cpu().numpy() # ndarray轉(zhuǎn)tensor t = torch.from_numpy(t).float()

PIL轉(zhuǎn)換

PyTorch中圖片采用[N,C,H,W]的順序存放，且數(shù)值在[0,1]之間，其他任何圖像庫讀取的圖片要在PyTorch中使用必須規(guī)范為該格式。這個轉(zhuǎn)換可以自己通過張量變換完成，也可以直接調(diào)用torchvision封裝好的函數(shù)。

import torch import torchvision import numpy as np from PIL import Image# tensor轉(zhuǎn)pil image t = torch.randn(32, 3, 224, 224) image = Image.fromarray(torch.clamp(t*255, min=0, max=255).byte().permute(1, 2, 0).cpu().numpy()) image = torchvision.transforms.functional.to_pil_image(t)# pil image轉(zhuǎn)tensor path = 'test.jpg' tensor = torch.from_numpy(np.asarray(Image.open(path).convert('RGB'))).permute(2, 0, 1).float() / 255 tensor = torchvision.transforms.functional.to_tensor(Image.open(path))

單元素張量值獲取

有時候?qū)τ趌oss這些值雖然結(jié)果是一個張量，但是這個張量其實就只含有一個值，這種張量通過item方法取出這個值。

import torcha = torch.randn(1) print(a) print(a.item())

模型操作

模型定義

主要的注意點就是繼承自nn.Module類且需要定義前向傳播運算，下面是個簡單的示例，具體可以查看我之前關(guān)于模型的文章。

import torch import torch.nn as nn import torch.nn.functional as Fclass Net(nn.Module):def __init__(self):super(Net, self).__init__()self.conv1 = nn.Conv2d(in_channels=3, out_channels=32, kernel_size=(3, 3))self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3)self.pool2 = nn.MaxPool2d(2, 2)self.fc1 = nn.Linear(64*54*54, 256)self.fc2 = nn.Linear(256, 128)self.fc3 = nn.Linear(128, 101)def forward(self, x):x = self.pool1(F.relu(self.conv1(x)))x = self.pool2(F.relu(self.conv2(x)))x = x.view(-1, 64*54*54)x = F.relu(self.fc1(x))x = F.relu(self.fc2(x))x = self.fc3(x)return x

模型總參數(shù)量

對所有層參數(shù)數(shù)目求和即可。

model = Net() num_parameters = sum(torch.numel(parameter) for parameter in model.parameters()) print(num_parameters)

模型參數(shù)查看

通過model.state_dict()或者model.named_parameters()查看所有可訓(xùn)練參數(shù)，包括父類參數(shù)。

params = list(model.named_parameters()) (name, param) = params[0] print(name) print(param.grad)

模型可視化

通過pytorchviz或者pytorch-summary兩個工具進行結(jié)構(gòu)可視化。

模型權(quán)重初始化

這方面我在之前的博客中提到過，一般區(qū)分不同的層進行不同的初始化，例如模型中定義如下的成員函數(shù)。尤其注意，model.modules()和model.children()都會返回模型的成員層，但是model.modules()會迭代地遍歷模型的所有子層，而model.children()只會遍歷模型下的一層。

def init_weights(self):for m in self.modules():if isinstance(m, nn.Conv2d):nn.init.xavier_normal_(m.weight.data)if m.bias is not None:m.bias.data.zero_()elif isinstance(m, nn.BatchNorm2d):m.weight.data.fill_(1)m.bias.data.zero_()elif isinstance(m, nn.Linear):nn.init.normal_(m.weight.data, 0, 0.01)m.bias.data.zero_()

提取模型某層

model.modules()會返回模型中所有模塊的迭代器，它能夠訪問到最內(nèi)層。

new_model = nn.Sequential() for layer in model.named_modules():if isinstance(layer[1],nn.Conv2d):new_model.add_module(layer[0],layer[1])

使用預(yù)訓(xùn)練模型參數(shù)

使用預(yù)訓(xùn)練模型一般是加載模型訓(xùn)練好的參數(shù)，大多時候訓(xùn)練到部署還會涉及到GPU模型參數(shù)加載到CPU模型上。

model.load_state_dict(torch.load('model.pkl'), strict=False) model.load_state_dict(torch.load('model.pth', map_location='cpu'))

數(shù)據(jù)處理

視頻基本信息

通過opencv-python得到視頻的信息。

import cv2 video = cv2.VideoCapture('v.mp4) height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT)) width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH)) num_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT)) fps = int(video.get(cv2.CAP_PROP_FPS)) video.release()

圖像常用預(yù)處理

各種處理方法的含義我在數(shù)據(jù)準備的博客中介紹過。

import torchvision train_transform = torchvision.transforms.Compose([torchvision.transforms.RandomResizedCrop(size=224,scale=(0.08, 1.0)),torchvision.transforms.RandomHorizontalFlip(),torchvision.transforms.ToTensor(),torchvision.transforms.Normalize(mean=(0.485, 0.456, 0.406),std=(0.229, 0.224, 0.225)), ]) val_transform = torchvision.transforms.Compose([torchvision.transforms.Resize(256),torchvision.transforms.CenterCrop(224),torchvision.transforms.ToTensor(),torchvision.transforms.Normalize(mean=(0.485, 0.456, 0.406),std=(0.229, 0.224, 0.225)), ])

分類模型訓(xùn)練及驗證

訓(xùn)練的大致流程代碼。

for epoch in range(epochs):# 訓(xùn)練集訓(xùn)練train_loss = 0.0correct = 0.0total = 0.0for step, data in enumerate(train_loader):x, y = dataout = model(x)loss = criterion(out, y)optimizer.zero_grad()loss.backward()optimizer.step()_, pred = torch.max(out.data, 1)total += y.size(0)correct += (pred == y).squeeze().sum().numpy()train_loss += loss.item()if step % 100 == 0:print("epoch", epoch, "step", step, "loss", loss.item())train_acc = correct / total# 驗證集驗證valid_loss = 0.0correct = 0.0total = 0.0for step, data in enumerate(valid_loader):model.eval()x, y = dataout = model(x)out.detach_()loss = criterion(out, y)_, pred = torch.max(out.data, 1)valid_loss += loss.item()total += y.size(0)correct += (pred == y).squeeze().sum().numpy()valid_acc = correct / totalscheduler.step(valid_loss)writer.add_scalars('loss', {'train_loss': train_loss, 'valid_loss': valid_loss}, epoch)writer.add_scalars('accuracy', {'train_acc': train_acc, 'valid_acc': valid_acc}, epoch)

自定義損失

PyTorch中損失也是繼承自nn.Module類，通過forward方法計算損失值。

from torch import nn class MyLoss(nn.Moudle):def __init__(self):super(MyLoss, self).__init__()def forward(self, x, y):loss = torch.mean((x - y) ** 2)return loss

梯度裁減

這是一個防止梯度爆炸的有效手段，但也會帶來一些問題。

torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=20)

其他

模型狀態(tài)

訓(xùn)練和驗證是對模型的要求是不一樣的，通過模型前向計算輸出前調(diào)整模型狀態(tài)。（通過model.train()和model.eval()進行調(diào)整。）

交叉熵損失

PyTorch中的交叉熵損失不需要經(jīng)過Softmax，因為其內(nèi)置了softmax計算。同時，也不需要進行onehot編碼，內(nèi)置了。

梯度置零

model.zero_grad()會把整個模型的參數(shù)的梯度都歸零, 而optimizer.zero_grad()只會把傳入其中的參數(shù)的梯度歸零。

反向傳播

loss.backward()進行反向傳播之前要用用optimizer.zero_grad()清除累積的梯度。

GPU數(shù)據(jù)IO

盡量減少CPU和GPU之間的數(shù)據(jù)交互，因為這種交互很費時間。如要獲得每個batch的loss 和acc，可以將它們累積在GPU中等整個epoch結(jié)束后一起傳輸?shù)紺PU（使用loss.cpu()方法）會比每個batch 都進行一次GPU到CPU的IO快很多。

補充說明

本文主要介紹PyTorch中常用的一些代碼段，都是比較實用的代碼，部分參考PyTorch CookBook。本文涉及到的所有代碼都可以在我的Github找到，歡迎star或者fork。

總結(jié)

以上是生活随笔為你收集整理的PyTorch-常用代码的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯，歡迎將生活随笔推薦給好友。