PyTorch機械学習モデルのGPUメモリ消費の要因を細かくトラックする。

# はじめに
　本記事では今更ながら、**この世で一番シンプルな機械学習モデル**（全結合1層）を使って**MNIST**の分類課題を行い、訓練の各過程でGPU上にどれくらいメモリが割り当てられるかを確認、解釈を行った。
　`pytorch_memlab`(あるいは`torch.cuda.memory_allocated(device)`)で表示される、PyTorchのテンソルによって占有されているGPUメモリ消費に関して**説明が可能**だったので、それを雑多にまとめ、検証している。

# 実験
> ### **条件**

PyTorch == 1.13.1
GPU : NVIDIA Tesla A30
CUDA version == 12.1

MNISTデータセットを一層の全結合層で分類。
最適化手法はSGD。
以下の`main.py`（これは`models.py`をimportして動くファイル)を`python main.py`で実行し、分類を行った。
今回は訓練1エポック目の一番最初のバッチで重みの更新が行われる所だけに注目し、この時GPU(1台)上でのメモリの立ち上がり方を追跡した。
```main.py
 1 # import library
 2 import torch
 3 import torch.nn as nn
 4 from torchvision.datasets import MNIST
 5 import torchvision.transforms as transforms
 6 from torch.utils.data import random_split
 7 from torch.utils.data import DataLoader
 8 import torch.optim as optim
 9
10 # import model
11 from models import *
12
13 def main():
14     # data prepare
15     mnist_train_data = MNIST('path/to/data', train=True, download=True, transform=transforms.ToTensor())
16     train_size = int(len(mnist_train_data)*0.8)
17     mnist_train_data, mnist_validation_data = random_split(mnist_train_data, [train_size, len(mnist_train_data) - train_size])
18     mnist_test_data = MNIST('path/to/data', train=False, transform=transforms.ToTensor())
19     train_loader = DataLoader(mnist_train_data, batch_size=100, shuffle=True)
20     valid_loader = DataLoader(mnist_validation_data, batch_size=100, shuffle=True)
21     test_loader = DataLoader(mnist_test_data, batch_size=100, shuffle=True)
22
23     # gpu setup
24     if torch.cuda.is_available():
25        gpu_id = 0
26        device = torch.device('cuda:{}'.format(gpu_id))
27
28    # training setup
29    image_size = 28*28
30    model = simple_net_mnist(image_size, 10).to(device)
31    optimizer = optim.SGD(model.parameters(), lr=0.01)
32    criterion = nn.CrossEntropyLoss()
33
34    # train
35    for epoch in range(100):
36        train(model, device, train_loader, optimizer, criterion, epoch, image_size)
37
38 # train function
39 def train(model, device, train_loader, optimizer, criterion, epoch, image_size):
40    model.train()
41    running_loss = 0.0
42    for batch_idx, data in enumerate(train_loader):
43        inputs, labels = data
44        inputs, labels = inputs.to(device), labels.to(device)
45        optimizer.zero_grad()
46        inputs = inputs.view(-1, image_size)
47        output = model(inputs)
48        loss = criterion(output, labels)
49        loss.backward()
50        optimizer.step()
51        running_loss += loss.item()
52        if batch_idx % 480 == 479:
53            print(f'[{epoch+1}, {batch_idx + 1:5d}] loss : {running_loss/2000:.3f}')
54            running_loss = 0.0
55
56 if __name__ == '__main__':
57    main()
```
```models.py
 1 import torch.nn as nn
 2 
 3 class simple_net_mnist(nn.Module):
 4    def __init__(self, input_dim, output_dim):
 5        super(simple_net_mnist, self).__init__()
 6        self.fc = nn.Linear(input_dim, output_dim)
 7   
 8    def forward(self, x):
 9         out = self.fc(x)
10         return out
```
# やったこと
　上記`main.py`, `models.py`の各行の前後でテンソルによって占有されているGPUメモリ消費の変動を確認。`pytorch_memlab`か`torch.cuda.memory_allocated(device)`を使えば良い。どこでどのくらいメモリが割り当てられているかを確認した。

# 結果
前後でGPU上での割り当てメモリが増えた行を全て列挙すると以下の通り。数値の単位はByte。

1. <font color="green">**main.py 30行目**</font> : `model = simple_net_mnist(~).to(device)`
0 --> 32,256 (+**32,256**)

2. <font color="green">**main.py 44行目**</font> : `inputs, labels = inputs.to(device), ~`
32,256 --> 347,136 (+**314,880**)

3. <font color="green">**main.py 47行目**</font> : `output = model(inputs)`
347,136 --> 351,232 (+**4,096**)

4. <font color="green">**main.py 48行目**</font> : `loss = criterion(output, labels)`
351,232 --> 356,352 (+**5,120**)

5. <font color="green">**main.py 49行目**</font> : `loss.backward()`
356,352 --> 384,000 (+**27,648**)

6. <font color="green">**models.py 9行目**</font> ： `out = self.fc(x)`
347,136 --> 351,232 (+**4,096**)

これより、最大で**384,000** Byteが割り当てられていたことがわかる。

# 解釈
　初めに、本モデル( `simple_net_mnist` )のデータの流れを簡単にまとめる。
モデルへの入力にはバッチサイズ(100) × 28 × 28の画像が入り、これを`main.py`46行目でreshapeし、バッチサイズ × 784に変換する。その後、全結合層（784 to 10）を経てバッチ内の各画像で10値の予測確率が出力される。この図を踏まえ、実際GPU上で確保されているメモリ量について解釈を行う。
![スクリーンショット 2023-08-16 17.04.58.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/680028/b8d88386-b652-8c91-d9c5-014fd7107964.png)


## 1. <font color="green">**main.py 30行目**</font> ( `model = simple_net_mnist(~).to(device)` ) : +**32,256**
　ここではモデル( `simple_net_mnist` )のパラメータ数分のメモリが確保されている。このモデルのパラメータ数は784 to 10の変換を行う全結合層であるため、重み＋バイアスで784 × 10 + 10で、floatは４Byteであるため、
```math
(784 \times 10 + 10)\times 4 = 31,760 
```
より**31,760** Byteが確保されていることがわかる。これは実際の確保分32,256 Byteより496 Byte少ないものの、実際の確保メモリの <font color="red">**98.5**</font> %の大きさ($31,760/32,256\approx0.985$より)で見積もることができている。

## 2. <font color="green">**main.py 44行目**</font> ( `inputs, labels = inputs.to(device), ~` ) : **+314,880**
　ここは単純にデータを転送しているだけ。入力データ = バッチサイズ(100) × 28 × 28 とラベル = バッチサイズ × 10 それぞれfloat, intでともに４Byteであるため、
```math
(100 \times 28 \times 28 + 100) \times 4 = 314,000
```
より**314,000**Byteが確保されていることがわかる。これは実際の確保分314,880より880 Byte少ないものの、実際の確保メモリの <font color="red">**99.7**</font> %の精度($314,000/314,880\approx0.997$より)で見積もることができている。

## 3. <font color="green">**main.py 47行目**</font> ( `output = model(inputs)` ) ＝ 6. <font color="green">**models.py 9行目**</font>( `out = self.fc(x)` ) (両者は同じ) : +**4,096**
　ここはforwardを流している。forwardの中間出力(今回は1層のため、ここでは単純にforwardを経て得られた出力値=バッチサイズ×10)を保存し、
```math
100 \times 10 \times 4 = 4,000
```
より**4,000** Byteが確保されていることがわかる。実際の確保分4,096より96 Byte少ないものの、実際の確保メモリの <font color="red">**97.7**</font> %の精度($4,000/4,096\approx0.977$より)で見積もることができている。

## 4. <font color="green">**main.py 48行目**</font> ( `loss = criterion(output, labels)` ) : +**5,120**
　ここはcross entropy誤差を求めている。既にforwardを経て得られたoutputの情報もlabelの情報もあるはずなので、更に追加で必要なメモリとは一体何なのだろうか。一つは1バッチ内の各画像のロスを入れる箱で、これはバッチサイズ × 4 = 400(Byte)だろう。
　また、cross entropyを求める式は、真の分布（今回だとone-hotの10個の要素からなる分布）を$p(x)$, 予測分布（今回だと10個の実数の要素からなる分布）を$q(x)$とするとcross entropy $H(p, q)$は
```math
H(p, q) = -\sum_x p(x) \log (q(x))
```
なる式で表される。したがって、以下二つの可能性が考えられる。

可能性1：元のラベルはintで入っていたが、cross entropyを求めるために内部でone-hotに変換され、その分のメモリが確保されている可能性がある。

可能性2：モデルから出力されるoutputsの各予測値に対してlog (q(x))を算出しなければならない。この分のメモリが中間出力とは別に確保されている可能性がある。

いずれにせよ、モデルから出力されるoutputsと同じメモリを最初の400（Byte）と足し合わせて確保する必要があり、
```math
400 + 100 \times 10 \times 4 = 4,400
```
より**4,400**Byteが確保されていることがわかる。実際の確保分5,120より720 Byte少ない。これより今回の見積もりは、実際の確保メモリの <font color="red">**85.9**</font> %程度($4,400/5,120\approx0.859$より)であることがわかる。（見積もりの精度が低いので、見落としや解釈間違いがありそう..）

## 5. <font color="green">**main.py 49行目**</font> ( `loss.backward()` ) : +**27,648**
　ここはbackwardを計算している。今回最適化アルゴリズムとしてSGDを使ったので、勾配の情報さえもとまれば重みの更新が可能となる。勾配の情報は元の重みとは別のコピー配列を用意して格納する必要があり、（コピーしないと**勾配**か**元の重みの値**のどちらかの情報が消えてしまう）モデルのパラメータ数分を新たに別で確保する。また、誤差が一度求まってしまえば中間出力の値の保持は不要であるため、その削除が行われる。これらを総合すると、
```math
(784 \times 10 + 10) \times 4 - 100 \times 10 \times 4 = 27,400
```
より**27,400** Byteが確保されたことになる。この見積もりは、実際の確保分27,648 Byteよりも248 Byte少ないものの、 実際の確保メモリの<font color="red">**99.1**</font> %の大きさ($27,400/27,648\approx0.991$より)で見積もれていることがわかる。

# 解釈まとめ
　ここまで、典型的な機械学習モデルで段階ごとにGPU上でどれ程のメモリが確保されているかを見積もった。その見積りの精度はほとんどの段階で <font color="red">**97 % ~ 99 %**</font> だったことから、少なくともSGDを使った学習では、今回行った解釈でそれ程大きく外してはいないはずである。今回の解釈はまとめると以下の4段階になる。

**段階0. 初期状態 : 0**
**段階1. モデルをロード : + モデルのパラメータ数** <font color="green">**(main.py 30行目)**</font>
**段階2. フォワード : + バッチサイズ × ( 入力データ + 中間出力(=中間層から出力された値, 最終層から出力された値outputs) + 誤差取得用のoutputs(?) )**　　<font color="green">**(main.py 44, 47, 48行目)**</font>
**段階3. バックワード : + モデルのパラメータ数(段階1.とは別)を確保 & 中間出力を削除** <font color="green">**(main.py 49行目)**</font>

以上をまとめると、MNISTのSGD学習で、今回の確認方法で表示された最大GPU割り当てメモリ/**Byte**は段階3時点での割り当てメモリ

**maxparam = ( モデルのパラメータ数 $\times$ 2 + バッチサイズ $\times$ (入力データ + outputs) ) $\times$ 4**

で近似できる事になる。ただしこの式は普遍的なものではなく、中間層の次元やモデルのサイズのバランス次第ではフォワードで最大GPU割り当てメモリとなる可能性もあるため、その場合には式の形も変わってくるだろう。

　今回の例だと、モデルのパラメータ数=$784\times 10 + 10$、バッチサイズ=$100$、入力データ=$784$, outputs=$10$より、
```math
{\rm max param} = ((784\times 10 + 10) \times 2 + 100 \times (784 + 10) ) * 4 = 380,400
```
となり、実際の割り当ては最大で$384,000$ Byteであったため$380,400/384,000 = 0.991$より <font color="red">**99.1**</font> % の精度で確保メモリの見積もりができていることがわかる。

# 検証
## 1. `nvidia-smi`コマンドで表示されるメモリ使用量に洞察を与えられないか？
　メモリ使用量を確認する際に手軽でよく登場するコマンドとして`nvidia-smi`(あるいは`nvitop`)が挙げられる。ただしこのコマンドは非テンソル占有メモリもまとめて表示され、今回の解釈をストレートには適用できない。
　そこでここでは、テンソル占有メモリに影響が及ぶハイパラだけを変えたときに占有メモリがどう変化するかを確認した。具体的には先ほどと全く同じ全結合一層のアーキテクチャ、データはMNIST、最適化手法はSGDで揃え、バッチサイズだけを様々に変えて`nvidia-smi`で表示される使用メモリ量/**MiB**をプロットした。その結果、以下のような図が得られた。

<img width="500" alt="スクリーンショット 2023-08-20 16.47.00.png" src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/680028/17c8f577-f0f6-7593-9413-e39fe87e92c7.png">

このプロットは凡そ線形になっており、傾きは3,212 Byte / 1 Batchだった。この傾きを**解釈まとめ**で考察した**maxparamの式**と紐付けると、傾き(Byte / 1 Batch)は "(入力データ + outputs）× 4" に相当することが分かる。この値を算出すると、
```math
(784 + 10) \times 4 = 3,176
```
となることから、限りなく`nvidia-smi`コマンドで表示されるメモリ使用量のバッチサイズに対する変化率3,212 Byte / 1 Batchを近似できており、**解釈まとめ**で示した**maxparamの式**の一定の妥当性が示唆された。

　なお、線形とは言いつつ切片の値が異常に高くなっていそうだという点については、CUDAの初期化時点で色々なカーネルをロードするCUDAコンテキストが作成されることによるものらしい。（ここら辺の非テンソル占有メモリについてはよくわからなかったため撤退した。他にも原因はありそうだが、一因と思しき[リンク](https://discuss.pytorch.org/t/confusion-on-tensors-memory-usage/162753)のdiscussionをご参照頂ければ幸いである。）

## 2-1. Adamではどうか？
　SGD以外の最適化アルゴリズムでも各行の前後で`torch.cuda.memory_allocated(device)` or `pytorch_memlab`によりメモリを確認した。コード`main.py`の31行目を
```main.py
~  ...
31 optimizer = optim.Adam(model.parameters(), lr=0.01)
~  ...
```
に置き換えて同じ実験を行った結果、

・GPU上で確保されているメモリはSGDの時と<font color="green">**main.py49行目**</font>までは同じ結果だった。

・ただし、<font color="green">**main.py50行目**</font>に記述されている`optimizer.step()`で384,000 --> 448,512へと+**64,512** Byte分のメモリ上昇が確認された。

## 2-2. Adam実験の解釈は？
　この上昇分はAdamで必要となる勾配情報の1次のモーメントと2次のモーメントのメモリ確保によるものだと考えられ、モデルのパラメータ数を2倍した式
```math
(784 \times 10 + 10) \times 4 \times 2 = 62,800
```
により、**62,800** Byteと見積もられた。この見積もりは実際に確保されたメモリ64,512 Byteの <font color="red">**97.3 %**</font> という精度であった。このことから、`loss.backward`時点まではSGDと**同じ挙動をしている**が、更新`optimizer.step()`の際に勾配情報に加えて新しく**1次・2次のモーメント用のメモリが割り当てられ**、これらを使って重みの更新が行われるということが示唆された。

# 結論
　今回確認、検証をしてみて、機械学習モデルのGPUへのメモリ割り当ては大きく以下の4段階で行われると解釈された。これらの解釈で見積もられたGPU割り当てメモリサイズは実際の割り当てメモリサイズの<font color="red">**97 % ~ 99 %**</font>であった。

・**段階1**：モデルをロード。
・**段階2**：フォワードでデータを流す。この時中間層からの出力の分もメモリを確保。
・**段階3**：バックワードで勾配を取得。そして中間出力の分のメモリは削除。
・**段階4**：重み更新。SGDでは新しいメモリ確保は不要。Adamなどの他の最適化アルゴリズムではモデルのパラメータ数×n個分のメモリを新しく確保。
　
　今回扱った1層の全結合モデルでSGDを用いてMNISTを分類する例では、GPU割り当てメモリは段階３で最大となり、既出のmaxparamの式で近似ができた。いつ割り当てメモリが最大となるかは**モデルと中間出力の次元の両方**により決まるため、それ次第では他の段階時点での割り当てメモリの近似式を用いるべきだと考えられる。

# 参考文献
・https://pytorch.org/docs/stable/cuda.html
・https://runebook.dev/ja/docs/pytorch/cuda?page=2
・https://medium.com/deep-learning-for-protein-design/a-comprehensive-guide-to-memory-usage-in-pytorch-b9b7c78031d3