ビデオメモリを賢く活用してください。ビデオメモリを節約するための 17 の pytorch トリック

ガイド

1. ビデオメモリはどこで使用されますか?

一般に、ニューラルネットワークをトレーニングする場合、ビデオメモリは主にネットワークモデルと中間変数によって占有されます。

ネットワークモデルの畳み込み層、全結合層、正規化層のパラメーターはビデオメモリを占有しますが、アクティベーション層とプーリング層は基本的にビデオメモリを占有しません。
中間変数には、グラフィックスメモリを最も消費する機能マップとオプティマイザーが含まれます。
実際、pytorch 自体もビデオメモリをある程度占有しますが、それほど多くはありません。おおよそ次の方法が優先順位として推奨されています。

ここに画像の説明を挿入します

2. ヒント 1: インプレース操作を使用する

インプレース演算（インプレース）とは文字通り、変数をその場で操作することを意味し、pytorchに相当するもので、新たなメモリ空間を適用せずに元のメモリ上で変数を操作し、メモリ使用量を削減することを意味します。具体的には、オンサイト運用には次の 3 つの実装側面が含まれます。

inplace 属性を True として定義するアクティベーション関数を使用します。たとえば、nn.ReLU(inplace=True)
インプレース操作で pytorch メソッドを使用します。通常、メソッド名の後には、、、などのアンダースコア "_" が続きtensor.add_()ますtensor.scatter_()。F.relu_()
、などの、その場で動作する演算子を使用しますy += x。y *= x

3. ヒント 2: 中間変数を避ける

カスタムネットワーク構造のメンバーメソッドの forward 関数では、不要な中間変数の使用を避け、以前に適用されたメモリで動作するようにしてください。たとえば、次のコードは中間変数を多用しすぎて、多くのメモリを消費します。不要なビデオメモリ:

def forward(self, x):

	x0 = self.conv0(x)  # 输入层
	x1 = F.relu_(self.conv1(x0) + x0)
	x2 = F.relu_(self.conv2(x1) + x1)
	x3 = F.relu_(self.conv3(x2) + x2)
	x4 = F.relu_(self.conv4(x3) + x3)
	x5 = F.relu_(self.conv5(x4) + x4)
	x6 = self.conv(x5)  # 输出层

	return x6

メモリ使用量を削減するために、上記の forward 関数を次のように変更できます。

def forward(self, x):

	x = self.conv0(x)  # 输入层
	x = F.relu_(self.conv1(x) + x)
	x = F.relu_(self.conv2(x) + x)
	x = F.relu_(self.conv3(x) + x)
	x = F.relu_(self.conv4(x) + x)
	x = F.relu_(self.conv5(x) + x)
	x = self.conv(x)  # 输出层

	return x

上記 2 つのコードで実装される機能は同じですが、ビデオメモリの使用方法が大きく異なり、後者は前者が占有するビデオメモリのほぼ 90% を節約できます。

4. ヒント 3: ネットワークモデルを最適化する

ネットワークモデルによるビデオメモリの占有は、主に畳み込み層、全結合層、正規化層のパラメータを指します。具体的な最適化方法には次のものが含まれますが、これらに限定されません。

コンボリューションカーネルの数を減らす (= 出力特徴マップチャネルの数を減らす)
完全に接続されたレイヤーを使用しないでください
nn.AdaptiveAvgPool2d()完全に接続されたレイヤーの代わりにグローバルプーリングnn.Linear()
正規化層を使用しない
ジャンプ接続のスパンは大きすぎてはなりません (多数の中間変数の生成を避けるため)。

5. ヒント 4: BATCH_SIZE を減らす

畳み込みニューラルネットワークをトレーニングする場合、エポックはデータ全体がトレーニングされる回数を表し、バッチはトレーニングに参加するためにエポックを batch_size のバッチに分割することを表します。
ビデオメモリの使用量を減らすには、batch_size を減らすのが一般的です。トレーニングが不十分な場合は、batch_size を減らすことが優先されます。ただし、batch_size を無限に小さくすることはできません。大きすぎるとネットワークが不安定になり、小さすぎると、batch_size を減らすことができます。、ネットワークは収束しません。

6. ヒント 5: バッチを分割する

バッチの分割は、ヒント 4 のバッチサイズを減らすこととは本質的に異なります。バッチを分割する操作は、2 つのトレーニングの損失を加算してから逆伝播することとして理解できますが、バッチサイズを減らす操作は、1 回トレーニングして逆伝播を実行することです。 . 一度広げます。バッチ分割操作は、元のバッチサイズを仮定すると、次の 3 つのステップとして理解できますbatch_size=64。

バッチを 2 つのbatch_size=32小さなバッチに分割します
ネットワークとターゲット値を個別に入力して損失を計算し、結果として生じる損失を加算します。
バックプロパゲーションを実行する

7. ヒント 6: PATCH_SIZE を小さくする

畳み込みニューラルネットワークのトレーニングでは、patch_size は入力ニューラルネットワークの画像サイズ、つまり (H*W) を指します。
ネットワーク入力パッチのサイズは、後続の特徴マップのサイズに大きな影響を与えます。トレーニング中に [64*64]、[128*128] などのサイズのパッチが使用される場合があります。ビデオメモリが不十分な場合、サイズは[32 * 32]、[16 * 16] など、パッチのサイズをさらに減らすこともできます。
ただし、この方法には問題があり、ネットワークの汎化能力に大きな影響を与える可能性があり、トリミングする際には元の画像をランダムにトリミングすることに注意する必要があり、一般的には推奨されません。

8. ヒント 7: 損失合計の最適化

バッチトレーニングの終了時に、対応する損失値が取得されます。エポックの損失を計算したい場合は、以前のすべてのバッチ損失を累積する必要があります。ただし、以前のバッチ損失は GPU のビデオメモリを占有します。直接累積して得られたエポックロスも GPU 内にあり、ビデオメモリを占有しますが、次の方法で最適化できます。

epoch_loss += batch_loss.detach().item()  # epoch 损失

上記のコードの効果は、まず、batch_loss テンソルの GPU 占有を解放し、テンソル内のデータを取り出して、それを蓄積することです。

9. ヒント 8: トレーニングの精度を調整する

トレーニング精度を低下させます。pytorch
でニューラルネットワークをトレーニングする場合、浮動小数点数はデフォルトで 32 ビット浮動小数点データを使用します。高い精度を必要としないネットワークをトレーニングする場合は、トレーニング用に 16 ビット浮動小数点データに変更できますが、データを組み合わせる場合は注意してください。ネットワークモデルは 16 ビット浮動小数点データに変換されます。そうしないと、エラーが報告されます。浮動小数点データを削減する実装プロセスは非常に簡単です。ただし、オプティマイザが Adam を選択した場合、エラーが報告される可能性があります。SGD オプティマイザが選択された場合、エラーは報告されません。具体的な操作手順は次のとおりです:

model.cuda().half()  # 网络模型设置半精度
# 网络输入和目标设置半精度
x, y = Variable(x).cuda().half(), Variable(y).cuda().half()

混合精度トレーニング混合
精度トレーニングとは、GPU を使用してネットワークをトレーニングするときに、計算を高速化するために関連データが半精度でメモリに保存および乗算され、丸め誤差を避けるために累積には全精度が使用されることを意味します。トレーニング時間は約半分に短縮され、メモリ使用量も大幅に削減できます。pytorch1.6 以前は、NVIDIA が提供する apex ライブラリを学習に使用することが主でしたが、それ以降は、pytorch に付属の amp ライブラリを使用することが多くなりました。

import torch
from torch.nn.functional import mse_loss
from torch.cuda.amp import autocast, GradScaler

EPOCH = 10  # 训练次数
LEARNING_RATE = 1e-3  # 学习率

x, y = torch.randn(3, 100).cuda(), torch.randn(3, 5).cuda()  # 定义网络输入输出
myNet = torch.nn.Linear(100, 5).cuda()  # 实例化网络，一个全连接层

optimizer = torch.optim.SGD(myNet.parameters(), lr=LEARNING_RATE)  # 定义优化器
scaler = GradScaler()  # 梯度缩放

for i in range(EPOCH):  # 训练

    with autocast():  # 设置混合精度运行
        y_pred = myNet(x)
        loss = mse_loss(y_pred, y)

    scaler.scale(loss).backward()  # 将张量乘以比例因子，反向传播
    scaler.step(optimizer)  # 将优化器的梯度张量除以比例因子。
    scaler.update()  # 更新比例因子

10. ヒント 9: トレーニングプロセスを分割する

非常に深いネットワークである resnet101 など、トレーニングされたネットワークが非常に深い場合、ディープニューラルネットワークを直接トレーニングするには非常に多くのグラフィックスメモリが必要となり、一般にネットワーク全体を一度に直接トレーニングすることはできません。この場合、複雑なネットワークを 2 つの小さなネットワークに分割し、別々にトレーニングすることができます。
Checkpoint は、時間をスペースと交換する pytorch のビデオメモリ不足に対するソリューションです。この方法は基本的に、トレーニングネットワーク全体に関係するパラメータの量を削減します。以下はコード例です。

import torch
import torch.nn as nn
from torch.utils.checkpoint import checkpoint

# 自定义函数
def conv(inplanes, outplanes, kernel_size, stride, padding):
    return nn.Sequential(nn.Conv2d(inplanes, outplanes, kernel_size, stride, padding),
                         nn.BatchNorm2d(outplanes),
                         nn.ReLU()
                         )


class Net(nn.Module):  # 自定义网络结构，分为三个子网络
    def __init__(self):
        super().__init__()

        self.conv0 = conv(3, 32, 3, 1, 1)
        self.conv1 = conv(32, 32, 3, 1, 1)
        self.conv2 = conv(32, 64, 3, 1, 1)
        self.conv3 = conv(64, 64, 3, 1, 1)
        self.conv4 = nn.Linear(64, 10)  # 全连接层

    def segment0(self, x):  # 子网络1
        x = self.conv0(x)
        return x

    def segment1(self, x):  # 子网络2
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        return x

    def segment2(self, x):  # 子网络3
        x = self.conv4(x)
        return x

    def forward(self, x):

        x = checkpoint(self.segment0, x)  # 使用 checkpoint
        x = checkpoint(self.segment1, x)
        x = checkpoint(self.segment2, x)

        return x

ネットワークトレーニングにチェックポイントを使用するには、属性を入力する必要がありますrequires_grad=True。指定されたコードでは、ネットワーク構造はトレーニングのために 3 つのサブネットワークに分割されます。ニューラルネットワークがnn.Sequential()構築されていない場合、それはさらにいくつかのカスタムサブネットワークにすぎません。例の Network ブロックは、のように個別に構築されます。
に含まれる大きなネットワークブロックの場合nn.Sequential()(小さなネットワークブロックには必要ありません)、checkpoint_sequentialパッケージを使用して実装を簡素化できます。具体的な実装プロセスは次のとおりです。

import torch
import torch.nn as nn
from torch.utils.checkpoint import checkpoint_sequential


class Net(nn.Module):  # 自定义网络结构，分为三个子网络
    def __init__(self):
        super().__init__()
		linear = [nn.Linear(10, 10) for _ in range(100)]
        self.conv = nn.Sequential(*linear)  # 网络主体，100 个全连接层

    def forward(self, x):

        num_segments = 2  # 拆分为两段
        x = checkpoint_sequential(self.conv, num_segments, x)

        return x

11. ヒント 10: ジャンクメモリをクリーンアップする

Python で定義された変数は、通常、使用後すぐにリソースを解放しません。次のコードを使用して、トレーニングサイクルの開始時にメモリのガベージをリサイクルできます。

import gc 
gc.collect()  # 清理内存

12. ヒント 11: 勾配累積を使用する

ビデオメモリサイズの制限により、大規模なネットワークモデルをトレーニングする場合は、より大きな batch_size を使用できません。一般に、batch_size を大きくすると、ネットワークモデルの収束が速くなります。
勾配累積は、複数のバッチによって計算された損失を平均してからバックプロパゲーションを実行することです。これは、ヒント 5 のバッチを分割するというアイデアに似ています (ただし、ヒント 5 は、大きなバッチを小さなバッチに分割することであり、トレーニングは依然として大きなバッチ、勾配の累積トレーニングは小さなバッチです)。
勾配累積の考え方を使用して、batch_size を大きくした場合の効果をシミュレートできます。具体的な実装コードは次のとおりです。

output = myNet(input_)  # 输入送入网络
loss = mse_loss(target, output)  # 计算损失
loss = loss / 4  # 累积 4 次梯度
loss.backward()  # 反向传播
if step % 4 == 0:  # 如果执行了 4 步
    optimizer.step()  # 更新网络参数
    optimizer.zero_grad()  # 优化器梯度清零

13. ヒント 12: 不要なグラデーションを削除する

テストプログラムの実行時には勾配関連の操作が含まれないため、ビデオメモリを節約するために、次の操作を含む不要な勾配をクリアできます。

コードを使用して、model.eval()正規化やニューロンのランダムな破棄などの操作を有効にせずにモデルをテスト状態にします。
テストコードはコンテキストマネージャーに配置されwith torch.no_grad():、グラフ構築などの操作は実行されません。
各トレーニングまたはテストサイクルの開始時に勾配クリア操作を追加します。

myNet.zero_grad()  # 模型参数梯度清零
optimizer.zero_grad()  # 优化器参数梯度清零

14. ヒント 13: ビデオメモリを定期的にクリーンアップする

同様に、pytorch 独自のコードを使用して、各トレーニングサイクルの開始時にビデオメモリをクリーンアップして、未使用のビデオメモリリソースを解放することもできます。

 torch.cuda.empty_cache()  # 释放显存

このステートメントの実行によって解放されたビデオメモリリソースは、Nvidia-smi コマンドで表示すると反映されませんが、実際には解放されています。実際、pytorch は原理的に変数が参照されなくなったら自動的に解放するので、この記述はあまり役に立たないかもしれませんが、個人的にはある程度役に立つと思います。

15. ヒント 14: ダウンサンプリングをもっと活用する

ダウンサンプリングは、実装の点ではプーリングに似ていますが、プーリングに限定されるものではなく、1 より大きいステップサイズを使用して、ダウンサンプリングのプーリングやその他の操作を置き換えることもできます。結果から判断すると、ダウンサンプリングによって得られる特徴マップが削減されます。特徴マップが削減されると、当然のことながらパラメータの数が削減され、ビデオメモリが節約されます。これは次の 2 つの方法で実現できます。

nn.Conv2d(32, 32, 3, 2, 1)  # 步长大于 1 下采样

nn.Conv2d(32, 32, 3, 1, 1)  # 卷积核接池化下采样
nn.MaxPool2d(2, 2)

16. ヒント 15: 無駄な変数を削除する

del 関数は、変数を完全に削除します。変数を再度使用する場合は、変数を再作成する必要があります。del は、メモリからデータを削除するのではなく、変数を削除することに注意してください。このデータは、他の変数によって参照される場合もあります。実装方法は次のように非常に簡単です。

 def forward(self, x):
 
	input_ = x
	x = F.relu_(self.conv1(x) + input_)
	x = F.relu_(self.conv2(x) + input_)
	x = F.relu_(self.conv3(x) + input_)
	
	del input_  # 删除变量 input_

	x = self.conv4(x)  # 输出层
	return x

17. ヒント 16: オプティマイザーを変更する

ネットワークトレーニングでより一般的に使用されるオプティマイザーは、SGD と Adam です。トレーニングの最終的な効果は別として、SGD は Adam よりもビデオメモリの使用量が少なくなります。他に方法がない場合は、パラメータ最適化アルゴリズムの変更を試みることができます。各最適化アルゴリズムは似ています。

import torch.optim as optim
from torchvision.models import resnet18

LEARNING_RATE = 1e-3  # 学习率
myNet = resnet18().cuda()  # 实例化网络

optimizer_adam = optim.Adam(myNet.parameters(), lr=LEAENING_RATE)  #  adam 网络参数优化算法
optimizer_sgd = optim.SGD(myNet.parameters(), lr=LEAENING_RATE)  # sgd 网络参数优化算法

18. 究極のヒント

十分なビデオメモリを搭載したグラフィックスカードを購入します。1 枚のカードが動作しない場合は、さらに数枚購入してください。

ここに画像の説明を挿入します