PyTorch 常用代碼整理

從網上看到別人整理的PyTorch常用代碼¹， 覺得很有用，所以就整理到自己的博客上，方便下次遺忘使用。
基於PyTorch1.0版本，需要用到以下包

import collections
import os
import shutil
import tqdm

import numpy as np
import PIL.Image
import torch
import torchvision

1. 基礎配置

檢查PyTorch版本

torch.__version__               # PyTorch version
torch.version.cuda              # Corresponding CUDA version
torch.backends.cudnn.version()  # Corresponding cuDNN version
torch.cuda.get_device_name(0)   # GPU type

更新PyTorch

PyTorch將安裝在anaconda3/lib/python3.7/site-packages/torch/目錄下（python2.x即將被淘汰，建議使用python3.x）

conda update pytorch torchvision -c pytorch

固定隨機種子

torch.manual_seed(0)
torch.cuda.manual_seed_all(0)

將程序運行在指定GPU顯卡上

在命令行下指定環境變量

CUDA_VISIBLE_DEVICES=0,1 python train.py  # notice has 'S' of 'DEVICES'

或者在代碼中指定

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

判斷是否有CUDA支持

torch.cuda.is_available()

設置爲cuDNN benchmark模式
Benchmark模式會提升計算速度，但計算中有隨機性，每次網絡前饋結果有差異。

torch.backends.cudnn.benchmark = True

如果想避免結果波動，添加設置

torch.backends.cudnn.deterministic = True

清除GPU存儲

有時Control-C終止程序會使得GPU存儲未釋放，需要手動清空。在PyTorch內部可以使用

torch.cuda.empty_cache()

或者在命令行先使用ps查找程序PID，再使用kill結束進程

ps aux | grep python
kill -9 [pid]

或者直接重置沒有清空的GPU（本人一般不用這種方式）

nvidia-smi --gpu-reset -i [gpu_id]

2. 張量處理

張量基本信息

tensor.type()    # Data type
tensor.size()    # Shape of the tensor. It is a subclass of Python tuple
tensor.dim()     # Number of dimensions.

數據類型轉換

# Set default tensor type. Float in PyTorch is much faster than double. 
torch.set_default_tensor_type(torch.FloatTensor)

# Type convertions.
tensor = tensor.cuda()
tensor = tensor.cpu()
tensor = tensor.float()
tensor = tensor.long()

torch.Tensor與np.ndarray轉換

# torch.Tensor -> np.ndarray.
ndarray = tensor.cpu().numpy()

# np.ndarray -> torch.Tensor
tensor = torch.from_numpy(ndarray).float()
tensor = torch.from_numpy(ndarray.copy()).float() # If ndarray has negative stride

torch.Tensor與PIL.Image轉換

PyTorch中的張量默認採用N*D*H*W的順序，並且數值範圍在[0,1]，需要進行轉置和規範化。

# torch.Tensor -> PIL.Image.
image = PIL.Image.fromarray(torch.clamp(tensor*255, min=0, max=255).byte().permute(1, 2, 0).cpu().numpy())
image = torchvision.transforms.functional.to_pil_image(tensor) # Equivalently way

# PIL.Image -> torch.Tensor
tensor = torch.from_numpy(np.asarray(PIL.Image.open(path))).permute(2, 0, 1).float() / 255
tensor = torchvision.transforms.functional.to_tensor(PIL.Image.open(path))  # Equivalently way

np.ndarray與PIL.Image轉換

# np.ndarray -> PIL.Image.
image = PIL.Image.fromarray(ndarray.astype(np.uint8))

# PIL.Image -> np.ndarray.
ndarray = np.asarray(PIL.Image.open(path))

從只包含一個元素的張量中提取值

這個在訓練時統計loss的變換過程中特別有用。否則這將累積計算圖，使GPU存儲佔用量越來越大。

value = tensor.item()

---

1. https://zhuanlan.zhihu.com/p/59205847 ↩︎