# Pytorch device-agnostic code

###### tags: `ML` `DL` `pytorch`

Pytorch 可以把每一個 tensor 都放在不同的 device 裡，所以如果要寫一份 code 可以在不同的 device 上面跑，必須要先做規劃，不然要換 device 的時候會變得非常麻煩。

[Pytorch doc](https://pytorch.org/docs/stable/notes/cuda.html#device-agnostic-code) 也有基本的介紹，這裡算是從不同角度切入來看，再從自己踩過的坑補充一些細節。

## 指定 device 的方法
- Tensor
- `torch.ones([2,5], device='cuda')`
- `torch.ones([2,5]).to('cuda')`
- `torch.ones_like(tensor)`（會得到一個和 `tensor` 一樣 shape, dtype, device 的 tensor）
- nn.Module
- `model.to('cuda')`

其中，`model.to` 的作法會把 model 裡的所有

- parameters
- buffers
- submodules

全部都換到指定的 device。

其實這 3 種幾乎已經包含所有 model 中會用到的 tensor 了。底下會說明怎麼利用這些特性，才能用最簡單的 `model.to` 一次把所有需要的 tensor 都換到正確的 device。

## Parameter & Buffer

如果在 model 裡需要紀錄自訂的 tensor，分為兩種做法：

- 如果是需要 autograd 更新的 [parameter](https://pytorch.org/docs/stable/nn.html#parameters)：`self.my_param = nn.Parameter(tensor)`
- 如果是不需要更新的 [buffer](https://pytorch.org/docs/stable/nn.html#torch.nn.Module.register_buffer)：`self.register_buffer('my_buffer', tensor)`
- register_buffer 之後，就可以用 `self.my_buffer` 拿到這個 tensor
- 這個 tensor 會被 model track，可以在 `model.state_dict()` 裡看到

這樣，這些 tensor 就都會在 `model.to` 的時候轉換到需要的 device。

需要特別注意的是 **buffer** 的部分，如果直接寫 `self.my_buffer = tensor`，那這個 tensor 就不會被 model track，也就是在 `model.state_dict()` 裡不會拿到，**同時，在 `model.to` 的時候也不會轉換 `self.my_buffer` 的 device**。

## Submodule

如果在 model 裡訂一個 attribute `self.my_module = my_module`，而且 `my_module` 是 nn.Module 的 instance（`isinstance(my_module, nn.Module)`），那 `my_module` 就是 model 的 submodule。

所以說如果你剛好想要一個 list 的 modules 作為 attribute：

- ❌ 如果直接寫 `self.my_modules = [my_module_1, my_module_2]`，那 `my_module_1` 和 `my_module_2` 就不是 model 的 submodule，在 `model.state_dict` 和 `model.parameters` 都不會出現。
- ✅ 應該用 [nn.ModuleList](https://pytorch.org/docs/stable/nn.html#torch.nn.ModuleList) 把這個 list 包起來：`self.my_modules = nn.ModuleList([my_module_1, my_module_2])`。

同理，如果想要一個 dict 的 modules 就用 `nn.ModuleDict`，如果要一個 list 的 parameters 就用 `nn.ParameterList`，要一個 dict 的 paramters 就用 `nn.ParameterDict`。

這樣這些 attributes 才會真正成為這個 model 的 submodule/parameter，支援各種功能，比如 `model.parameters()` 會 iterate over 所有 paramters 和 submodule 中的 parameters；`model.state_dict()` 會包含所有 submodule 中的 parameters 和 buffers，然後 `model.to('cuda')` 的時候也才會把全部 submodule 中的 tensors 都轉成 cuda tensor。

## Forward 時產生 Tensor

如果在 forward 時需要生成新的 tensor，比如 VAE 裡的 random noise

https://github.com/pytorch/examples/blob/5b1f45057dc14a5e2132b45233c258a1dc2a0aab/vae/main.py#L55

- ❌ 如果直接寫 `eps = torch.randn(*std.shape)`，那這個 eps 就會放在 cpu 裡，整個 model 就沒辦法在其他 device 上跑了。
- ❓ 雖然可以用 global variable 或者 model attribute 判斷現在使用的 device （`eps = torch.randn(*std.shape, device=device)` ），但是這樣等於是要在 `model.to` 以外的地方另外 maintain 一個自己的 device。
- ✅ 所以最好的做法就是用 `torch.randn_like`，這樣就可以確保這個 tensor 的 device 一定會和 input 相同。
- ✅ 如果因為 shape 不同之類的原因不能用 xxx_like，另外一個確保新 tensor 的 device 和 input 相同的方法：`torch.randn(*std.shape, device=std.device)`

## Input/Target tensor

Input/target tensor 做 device 轉換的時機視情況有所不同。

- 如果 dataset 很小，可以全部 load 到 gpu 的 memory 裡，那就在 dataset initialize 的時候直接把 tensor 都轉成 cuda tensor 會是最快的。
- 如果 dataset 比較大，必須一個 batch 一個 batch 的 load 到 gpu 裡，那就
1. 在 tensor 進 model 前做 device 轉換，並且把 DataLoader 的 pin_memory 設成 true（[討論](https://discuss.pytorch.org/t/should-we-set-non-blocking-to-true/38234/4)），但是這樣就是每次有任何 tensor 需要進 model 前都要 to 一次，要在 train 和 validation 都做一次也是有點麻煩。
2. 自訂 DataLoader 的 collate function，在 collate function 的最後做 device 轉換。這樣就是 model 只能吃 DataLoader 吐出來的 tensor，不然就又回到 (1)。
3. 在 model forward 裡用 `input.to(next(self.parameters()).device)` 做轉換。這是終極解決方案，只要 model 設定好 device 任何 tensor input model 都會進到這個 device。但是似乎很少看到有人這樣做，不知道為什麼？