# PyTorch 效能懶人包
[TOC]
## 1. 減少 I/O 時間
### 盡量不要從硬碟讀，能放 RAM 就放 RAM
Slow :-1:
```python
class Dataset(torch.utils.data.Dataset):
def __init__(self, data_root):
self.image_paths = glob.glob(os.path.join(data_root, '*.png'))
def __len__(self):
return len(self.image_paths)
def __getitem__(self, index):
return Image.open(self.image_path[index]).copy()
```
Fast :+1:
```python
class Dataset(torch.utils.data.Dataset):
def __init__(self, data_root):
image_paths = glob.glob(os.path.join(data_root, '*.png'))
self.images = [Image.open(path).copy() for path in image_paths]
def __len__(self):
return len(self.images)
def __getitem__(self, index):
return self.images[index]
```
- 只用 `Image.open(...)` 不會把圖片檔案關掉，開太多系統會爆炸，所以要記得 `.copy()`
- 因為加上 `copy()` 會 dereference `ImageFile` 物件，去 call destructor 把檔案關掉
- 記得估計 dataset 放進 RAM 會不會塞爆
- 如果 dataset 很大會塞爆，要確認 swap 夠大，放進 RAM 的大小如果超過 RAM+swap 整台 server 會卡到需要重開，如果是樓上的 server 如 Salmonella 和 COVID-19 還要申請才能上去重開。
- 看 swap 大小及用量可用 `htop`，詳細紀錄在 slack 的 server-maintain channal
- 部分從硬碟讀也是個策略
### LRU Cache
假如：資料超大，不想把所有圖片都讀進 RAM，但是有兩個 patch 是用同一張圖挖出來的，又不想讀兩次圖（你的 dataset class 可能會即時讀圖，因為全部讀進來太大了）
```python
from functools import lru_cache
@lru_cache(1000)
def read_image(path):
return Image.open(path).copy()
```
### 如果用到 DataLoader 記得設 `pin_memory=True`
好慢 :-1:
```python
loader = DataLoader(dataset)
for xs, ys in iter(loader):
xs, ys = xs.cuda(), ys.cuda()
...
```
好快 :+1:
```python
loader = DataLoader(dataset, pin_memory=True)
for xs, ys in iter(loader):
xs, ys = xs.cuda(), ys.cuda()
...
```
- 有用 CUDA 就設成 True。這個選項代表 使用 page-locked memory，讓 GPU 可以直接拿，不然 GPU 取得 tensor 資料時還要再自己 allocate 暫時性的 page-locked memory，會很慢。
https://devblogs.nvidia.com/how-optimize-data-transfers-cuda-cc/
> Host (CPU) data allocations are pageable by default. The GPU cannot access data directly from pageable host memory, so when a data transfer from pageable host memory to device memory is invoked, the CUDA driver must first allocate a temporary page-locked, or “pinned”, host array, copy the host data to the pinned array, and then transfer the data from the pinned array to device memory
- `tensor.cuda(non_blocking=True)` 可以異步 allocate 空間，有機會加速
https://stackoverflow.com/questions/55563376/pytorch-how-does-pin-memory-works-in-dataloader
### `/dev/shm`
這個目錄底下是掛在記憶體內的檔案系統，所以讀寫比硬碟快很多，如果記憶體夠大，你可以把 dataset 目錄整個放進去 `/dev/shm`。
```shell
mkdir /dev/shm/dataset
cp dataset/*.png /dev/shm/dataset
ln -s /dev/shm/dataset dataset_in_memory
```
## 2. 減少 CPU 運算時間
### DataLoader workers
- 如果 `__getitem__` 有處理資料，DataLoader 可開多點 worker
- 如果沒有就不要，因為建立 worker 也會花時間
- 看 CPU worker 數量可以用 `nproc`
```python
class Dataset(torch.utils.data.Dataset):
...
def __getitem__(self, index):
image = self.images[index]
image = crop_image(image)
image = rotate_image(image) # very CPU-bound if angle % 90 != 0
image = resize_image(image)
image = normalize_image(image)
image = binarize_image(image)
return image
dataset = Dataset('./data')
```
超慢 :-1:
```python
loader = torch.utils.data.DataLoader(dataset)
```
好快 :+1:
```python
loader = torch.utils.data.DataLoader(dataset, num_workers=48)
```
### 多用 torch.tensor 或者 np.array 來操作資料
- Python 原生物件慢到不行，把資料轉成 `np.ndarray` 再做運算，跟用 C 寫的速度差不多！
- PyTorch 的 CPU tensor 跟 Numpy ndarray 速度差不多
- 如果一次對多個值操作，`np.ndarray` 的 function 都很快，能用就用吧
```python
py_list = list(range(1000))
np_array = np.arange(1000)
max(py_list) # 14.9 µs
np_array.max() # 2.88 µs
# ~5x speedup
sum(a * b for a, b in zip(py_list, py_list)) # 88.3 µs
np.dot(np_array, np_array) # 1.89 µs
# ~47x speedup
```
### 東西先處理好存起來，訓練的時候不要浪費 CPU 資源
可以用 pickle 存任意 Python object
```python
import pickle
data = preprocess_dataset()
with open('data.pkl', 'wb') as f:
pickle.dump(data, f)
with open('data.pkl', 'rb') as f:
data = pickle.load(f)
```
可以用 numpy.save 存 np.array
```python
import numpy as np
data = np.random.randn(100, 100)
np.save('data.npy', data)
data = np.load('data.npy')
```
還有其他類似的東西可以快速存 list, dict, str 之類的物件
- torch.save / torch.load
- bson
- cbor
- msgpack
- 效能比較：
https://medium.com/@shmulikamar/python-serialization-benchmarks-8e5bb700530b
## 3. 增加 GPU 運算效率
### Batch 塞大坨一點
- GPU 就是拿來平行運算的
- 有時候 batch=32 跟 batch=64 跑的時間根本差不多，可以自己試試看
- 只要 memory 大小允許，batch 就設大一點，gradient 也會比較穩定
### 一張卡跑多個 model
- 在記憶體夠的情況下，如果 GPU util 很低
- 多跑幾個 model 比較不浪費 GPU 的資源，提高 GPU util
### CUDNN
```python
from torch.backends import cudnn
cudnn.benchmark = True # fast training
```
### 用 apex.amp 自動將 model 轉換成混合浮點精度
NVIDIA 的套件，幫你把模型轉成混合精度給 CUDA 算：
https://github.com/NVIDIA/apex
AMP = Automatic Mixed Precision
#### 安裝
1. 安裝時的環境必須要有 nvcc 和 g++，他要現場編譯。
2. 順便安裝 `libgomp`
3. Apex 會用到 PyTorch 的 package，因此 PyTorch 的 CUDA 版本要和系統的一樣，你可以用 `nvidia-smi` 來看現在的 driver 版本，如果 PyTorch 會動就是那個版本。
4. 如果編譯時出現版本不符，要下 `export CUDA_HOME=/usr/local/cuda-10.2` 和 `export PATH="/usr/local/cuda-10.2/bin:$PATH"`。`10.2` 可要改成你的 cuda 版本，如果沒有這個資料夾應該是因為你沒有 `sudo apt install cuda-toolkit-102`，`102` 就是 `10.2` 的意思。
```bash
mkdir -p /tmp/$USER
cd /tmp/$USER
git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./
```
- 備註：PyTorch 1.5.0 有自己做的 amp：`torch.cuda.amp`
https://pytorch.org/docs/stable/amp.html
https://pytorch.org/docs/stable/notes/amp_examples.html#gradient-scaling-examples
#### 片語
```python
# import
from apex import amp
# initialize model and optimizer
model, optimizer = amp.initialize(model, optimizer, opt_level='O1')
model, optimizer = amp.initialize(model, optimizer, opt_level='O2')
[model1, model2], [optimizer1, optimizer2] = amp.initialize([model1, model2], [optimizer1, optimizer2])
# rescale loss when backward
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
```
- Loss scaling 可讓過小的值不至於在低精度時消失。
https://docs.nvidia.com/deeplearning/performance/mixed-precision-training/index.html
- Advanced Usage: https://nvidia.github.io/apex/advanced.html
#### 重點
- 一定要記得 `scale_loss` 不然會 train 爛
- `opt_level` 有四種 `O0`、`O1`、`O2`、`O3`
- `O0` 沒任何效果，`O1`、`O2` 差不多，`O3` 很常爛掉（例如出現 NaN）
- PyTorch 偵測 NaN：
https://pytorch.org/docs/stable/autograd.html#torch.autograd.detect_anomaly
或者
```python
assert not torch.isnan(grad).any(), 'gradient contains NaN'
```
- 通常只會用 `O1` 和 `O2`，兩者只是實作上的差異而已。一般來說建議用 `O1`。
- Model 越大越複雜，或者 batch 或 input 越大，加速效果越明顯。（反過來說，小的 model 用 apex 有可能沒用或者變慢）
#### 範例
`# IMPORTANT PART` 是用 apex 要加的地方。
```python
import torch
from torch.optim import Adam
from torch.utils.data import DataLoader
from torchvision.models.densenet import densenet121
from torchvision.datasets import CIFAR10
import torchvision.transforms as transforms
# IMPORTANT PART
from apex import amp
# IMPORTANT PART
torch.cuda.set_device(0) # see: https://github.com/NVIDIA/apex/issues/319
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
dataset = CIFAR10(root='/tmp/cifar10', train=True, download=True, transform=transform)
loader = DataLoader(dataset, batch_size=128, shuffle=True, num_workers=16, pin_memory=True)
net = densenet121(num_classes=len(dataset.classes)).train().cuda()
opt = Adam(net.parameters())
# IMPORTANT PART
net, opt = amp.initialize(net, opt) # default is O1
loss_fn = torch.nn.CrossEntropyLoss()
for xs, ys in iter(loader):
xs = xs.cuda()
ys = ys.cuda()
logits = net(xs)
loss = loss_fn(logits, ys)
print('loss:', loss.item())
opt.zero_grad()
# IMPORTANT PART
with amp.scale_loss(loss, opt) as scaled_loss:
scaled_loss.backward()
opt.step()
```
#### 儲存與載入
- 轉成 apex 模型之後的 `state_dict` 跟原本不相容
- Checkpoint 也要順便存 amp state！
- Save:
```python
net = Net().cuda()
opt = Opt(net.parameters())
net, opt = amp.initialize(net, opt, opt_level='O1')
torch.save({
'net': net.state_dict(),
'opt': opt.state_dict(),
'amp': amp.state_dict(), # IMPORTANT PART
}, 'checkpoint.pt')
```
- Load:
```python
net = Net().cuda()
opt = Opt(net.parameters())
# using the same opt_level is recommended!
net, opt = amp.initialize(net, opt, opt_level='O1')
state = torch.load('checkpoint.pt')
net.load_state_dict(state['net'])
opt.load_state_dict(state['opt'])
amp.load_state_dict(state['amp']) # IMPORTANT PART
```
## 4. 其他
### `nn.DataParallel`
- 可以用多張 GPU 同時計算，一次跑大一點的 batch size
- https://pytorch.org/tutorials/beginner/blitz/data_parallel_tutorial.html
### Multiprocessing
- 訓練 Hogwild, A3C... 時很重要
- https://pytorch.org/docs/stable/notes/multiprocessing.html
### Distributed
- 一次用多台 server 跑
- https://pytorch.org/tutorials/intermediate/dist_tuto.html
### `torch.utils.bottleneck`
- https://pytorch.org/docs/stable/bottleneck.html
### PyTorch JIT (TorchScript)
- 把一系列運算編譯成靜態圖，可以存起來丟給其其他 frontend 跑。
- 存成計算圖的好處是我們可以一次丟比較多運算給 CUDA 跑，這會讓運算能比較好的被最佳化。
- 這篇提到可以用 TorchScript 自己寫 RNN Cell：
https://pytorch.org/blog/optimizing-cuda-rnn-with-torchscript/
### TensorFlow VS PyTorch VS MXNet
https://medium.com/syncedreview/tensorflow-pytorch-or-mxnet-a-comprehensive-evaluation-on-nlp-cv-tasks-with-titan-rtx-cdf816fc3935
Sign in with Wallet
Connect another wallet