NumPy/PyTorch

【極めればこのテンソル操作】インプレース操作でメモリ効率化！

Qualiteg プロダクト開発部

2024年10月22日 — 13 min read

Photo by Mohamed M / Unsplash

こんにちは！今日は PyTorchのインプレース操作に関する内容です！

ディープラーニングの学習モデルを作ってると、メモリ管理が大きな課題になります。課題の大部分はGPUメモリとお考えの方も多いのではないでしょうか。

そんなときに助けてくれるのが、PyTorchのインプレース操作です！

この記事では、インプレース操作の使い方をいろんな角度から見ていきたいとおもいます。

インプレース操作って何？

基本的な考え方

インプレース操作とは、既存のメモリ領域を直接書き換える操作のことです。PyTorchでは、演算子の後ろにアンダースコア（_）をつけることでインプレース操作を実行できます。

つまり、普通の操作だと新しいメモリを確保する必要がありますが、インプレース操作なら既存のメモリを直接書き換えることが可能です。

それでは、実際に見てみましょう！

import torch

# 普通の操作
x = torch.tensor([1, 2, 3])
y = x + 5  # 新しいメモリが必要

# インプレース操作ならこう！
x = torch.tensor([1, 2, 3])
x.add_(5)  # xを直接書き換え

さて、上記コードにおけるメモリの使われ方と操作の違いを詳しく見ていきましょう！

まず、通常の操作（x + 5）の場合を見てみましょう。この操作では、メモリ上に[1, 2, 3]というデータを持つテンソルxが作られた後、x + 5という操作が実行されると新しいメモリ領域が確保されます。

この新しい領域に[6, 7, 8]という計算結果が書き込まれ、その新しいメモリ領域への参照がyに代入されます。

この時、元のxは変更されず[1, 2, 3]のまま残ります。結果として、メモリ上にはxとyという2つの異なるテンソルが存在することになります。

一方、インプレース操作（x.add_(5)）では、処理の仕方が大きく異なります。まずメモリ上に[1, 2, 3]というデータを持つテンソルxが作られますが、add_メソッドが呼ばれると、xが使用している同じメモリ領域上で直接計算が行われます。
元のデータ[1, 2, 3]は[6, 7, 8]で上書きされ、新しいメモリ領域は一切確保されません。xは単に更新された値を指すようになります。

以下のように実際のメモリアドレスを確認すと違いが明確になります

# 通常の操作の場合
x = torch.tensor([1, 2, 3])
print(f"xのメモリアドレス: {x.data_ptr()}")
y = x + 5
print(f"yのメモリアドレス: {y.data_ptr()}")  # xとは異なるアドレス

# インプレース操作の場合
x = torch.tensor([1, 2, 3])
print(f"操作前のxのアドレス: {x.data_ptr()}")
x.add_(5)
print(f"操作後のxのアドレス: {x.data_ptr()}")  # アドレスは同じ

この違いの重要性は、特に大きなテンソルを扱う際に顕著になります。

例えば1000×1000の行列に対する操作を考えてみましょう

# 1000×1000の行列の場合（float32型）
large_x = torch.randn(1000, 1000)  # 約4MB

通常の操作では、この4MBのデータに対して演算を行うたびに新たに4MBのメモリが必要になります。一方、インプレース操作では追加のメモリは必要なく、既存の4MBのメモリ領域を再利用することができます。

このメモリ効率の違いは、大規模なニューラルネットワークの推論時や、メモリに制約のある環境での実行時に特に重要になってきます。

さて、次は、より実践的な活用シーンでみていきましょう！

実践的な使い方

1. 推論時の最適化

以下のコードは、シンプルながら効率的なネットワークの実装例です

class EfficientNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, 64, 3)
        self.bn = nn.BatchNorm2d(64)
        
    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x.relu_()  # インプレースでReLU適用！
        return x

コードのポイントを詳しく見ていきましょう！

まず、モデルの構造について説明すると、入力として3チャンネル（RGBカラー画像を想定）を受け取り、64個のフィルターを持つ3×3の畳み込み層を通します。

その後、バッチ正規化を行って、最後にReLU活性化関数を適用する、という流れになっています。

特に注目して欲しいのは、ReLUの適用方法です。

普通ならx = torch.relu(x)やF.relu(x)と書くところを、x.relu_()というインプレース操作を使っています。

これには大きな利点があります

通常のReLU操作では、以下のようなメモリの動きが発生します

畳み込み層の出力用メモリ
バッチ正規化の出力用メモリ
ReLU用の新しいメモリ

でも、インプレース操作を使うと、ReLU用の新しいメモリは必要ありません。

バッチ正規化の出力を直接書き換えてしまうからです。

特に推論時は、バッチ正規化の出力を保持しておく必要がないので、このような最適化が可能になります。

画像処理のような大きなテンソルを扱う場合、この差は実はかなり大きなものとなります。

例えばvggやresnetなどで224×224のRGB画像をバッチサイズ32で処理する場合、この層の出力だけでも：

32 (バッチサイズ) * 64 (チャンネル数) * 222 * 222 (出力サイズ) * 4 (float32のバイト数)

というサイズのメモリが節約できます。

なお、学習時にこのコードを使う場合は注意が必要です。

バックプロパゲーションでReLUの勾配が必要になる場合は、インプレース操作を避けた方が安全です。

その場合は以下のように書き換えましょう。

def forward(self, x):
    x = self.conv(x)
    x = self.bn(x)
    x = F.relu(x)  # 学習時は通常のReLUを使用
    return x

2. データ前処理の効率化

データの前処理で最もよく使う操作の1つが正規化です。特に大規模なデータセットを扱う場合、この処理のメモリ効率は非常に重要になってきます。

以下のコードは、インプレース操作を活用した効率的な正規化の実装例となります

def efficient_normalize(tensor):
    mean = tensor.mean()
    std = tensor.std()
    tensor.sub_(mean).div_(std)  # チェーン化できます！
    return tensor

このシンプルな実装には、いくつかの工夫が含まれています。まず、mean()とstd()はスカラー値を返すため、メモリ使用量はほとんど無視できますね。
ここで重要なのは、正規化の計算をインプレース操作で行っている点です。

普通に書くと以下のようになりがちな処理を、

normalized = (tensor - mean) / std  # 新しいメモリ領域が必要

インプレース操作を使うことで、余分なメモリ確保を避けています。

さらに、sub_()とdiv_()をチェーンすることで、コードもスッキリまとまっています。

例えば、バッチサイズ32の画像データ（224×224のRGB画像）を正規化する場合は、

# バッチサイズ * チャンネル数 * 高さ * 幅 * 4バイト
32 * 3 * 224 * 224 * 4 = 約19MB

のように、このサイズのメモリが節約できることになります

データローダーで毎回の読み込み時に正規化を行う場合、この差も大きいですよね。

もちろん、元のデータを保持しておきたい場合は、以下のように使います。

normalized_data = efficient_normalize(original_data.clone())

3. 学習時の最適化処理での活用

以下のコードは、インプレース操作を活用したSGD（確率的勾配降下法）の実装例です。重みの更新は頻繁に使いますので、おさえておきたいですね。

def custom_sgd_update(parameters, lr):
    with torch.no_grad():
        for param in parameters:
            if param.grad is not None:
                param.sub_(lr * param.grad)  # インプレースで重み更新

with torch.no_grad():というコンテキストマネージャを使用することで、自動微分の追跡を無効化しています。

インプレース操作はどこにあるでしょうか。

はい、重みの更新にインプレース操作（sub_）を使用していますね。

通常の減算を使用した場合：

param = param - lr * param.grad  # 新しいメモリ領域が必要

このように、更新のたびに新しいメモリ領域が必要になってしまいます。

一方、インプレース操作を使用することで、

新しいメモリ割り当てが不要
メモリの解放・再割り当ての処理が省略可能
大規模なモデルでも効率的な更新が可能

となります。

例えば、5000万パラメータのモデルの場合は以下のサイズのメモリ領域を、更新のたびに確保・解放する必要がなくなります。

# パラメータ数 * 4バイト（float32の場合）
50,000,000 * 4 = 約190MB

こんな感じで使用します

# モデルのパラメータを更新
learning_rate = 0.01
custom_sgd_update(model.parameters(), learning_rate)

気をつけるべきポイント

自動微分との相性

PyTorchの自動微分システムとインプレース操作の関係について、具体的なコードで見ていきましょう

x = torch.tensor([1., 2., 3.], requires_grad=True)
y = x * 2

# これはOK！
z = y.relu()

# これは危険！
y.relu_()  # 勾配計算に必要な情報が消えちゃう...

このコードを通じて、インプレース操作の注意点を詳しく見ていきましょう。まず、xというテンソルを作成し、requires_grad=Trueを指定することで勾配計算を有効にしています。次に、このxに対して2倍の演算を行い、その結果をyに格納しています。

ここで重要なのは、PyTorchの自動微分システムが計算グラフを構築している点です。通常のReLU関数を使用する場合（z = y.relu()）、新しいテンソルzが作成され、元の値yは計算グラフ内に保持されます。これにより、バックプロパゲーション時に適切に勾配を計算することができます。

しかし、インプレース操作のrelu_()を使用すると、元の値yが直接書き換えられてしまいます。この場合、バックプロパゲーション時に必要な情報（活性化関数適用前の値）が失われてしまうため、正しい勾配計算ができなくなります。

具体的には、ReLUの勾配計算において、入力が正だったのか負だったのかを判断するための情報が消失してしまうのです。

実際にPyTorchは、このような危険な操作に対して実行時エラーを発生させることがあります。これは、計算グラフの一貫性を保護するための重要な安全機構です。直感的には、元の値を保持したまま新しい結果を生成する通常の操作は「履歴を残す」操作、インプレース操作は「履歴を上書きする」操作と考えることができます。

つまり、インプレース操作を学習時に使うのは危険？どっち？

学習時の順伝播では通常の操作を使用し、インプレース操作は推論時や計算グラフの構築が完了した後の重み更新フェーズなど、勾配計算に影響を与えない場面に限定して使用するのはアリです。

さきほど「最適化処理の実装」で重み更新でインプレース操作を例示したのに、前節では、「学習時の順伝播では通常の操作」とかいていて混乱してしまう読者の方がいるのではないでしゅか？

実はこの点はしっかり理解しておくべきポイントがあり、そこを少し掘り下げてみます。

学習時のインプレース操作：順伝播と重み更新の違い

ここでは、学習プロセスにおけるインプレース操作の使用について、混乱を招きやすい部分を整理してみます。

実は、「学習時」といっても、そのタイミングによってインプレース操作の使用可否が変わってきます。

学習の1イテレーションの流れ

まず、学習の1イテレーションの流れを見てみましょう：

# 1. 順伝播（forward pass）
output = model(input_data)         # ここではインプレース操作は避ける
loss = criterion(output, target)   # 損失計算

# 2. 逆伝播（backward pass）
loss.backward()                    # 勾配計算

# 3. 重み更新（parameter update）
optimizer.step()                   # ここではインプレース操作OK

この3ステップにおいて、インプレース操作の使用可否は以下のように分かれます：

順伝播では以下のように書くのは危険です

def forward(self, x):
    x = self.linear1(x)
    x.relu_()  # 危険！このタイミングでのインプレース操作は避ける
    return x

一方、重み更新時には以下のようなインプレース操作は安全です

def update_weights(parameters, lr):
    with torch.no_grad():
        for param in parameters:
            if param.grad is not None:
                param.sub_(lr * param.grad)  # OK！このタイミングならインプレース操作可能

この違いが生じる理由は、計算グラフとバックプロパゲーションの関係にあります

【順伝播時（危険な理由）】

この時点では、バックプロパゲーションのために計算グラフを構築している最中です
中間の計算結果は、後の勾配計算で必要になります
インプレース操作で値を書き換えてしまうと、勾配計算に必要な情報が失われてしまいます

【重み更新時（安全な理由）】

この時点では、すでにバックプロパゲーションが完了しています
現在の計算グラフでの計算はすべて終わっています
次のイテレーションでは新しい計算グラフが作られます
パラメータを書き換えても、現在の勾配計算には影響しません

つまり、「学習時はインプレース操作を避ける」というのは、より正確には「学習時の順伝播における中間計算ではインプレース操作を避ける」という意味です。一方で、「すでに勾配計算が完了した後の重み更新では、インプレース操作を使用できる」というわけです。

このように、インプレース操作の使用可否は、単に「学習時か否か」ではなく、「計算グラフの構築中か、すでに計算が完了しているか」という観点で判断する必要があります。

さらに、理解を深めるために、「最適化処理の実装」のコードをもう一度みてみましょう。

def custom_sgd_update(parameters, lr):
    with torch.no_grad():
        for param in parameters:
            if param.grad is not None:
                param.sub_(lr * param.grad)  # インプレースで重み更新

このコードでインプレース操作が安全に使える理由について、もういちどおさらいしましょう。

まず、この処理はwith torch.no_grad()のコンテキスト内で実行されています。これは自動微分の追跡を無効化するもので、重み更新フェーズではすでにバックプロパゲーションが完了しているため、新しい計算グラフを作る必要がないことを示しています。

このタイミングでは現在のバッチにおける勾配計算がすべて完了していること、が重要な事実です。

現在のパラメータや中間結果を参照する必要はもうなく、次のバッチでは新しい計算グラフが作られることになります。故に、このタイミングでインプレース操作を使用しても安全なのです。

一方で、学習時の順伝播での使用には注意が必要です。例えば以下のような実装は危険です：

# 危険な例！
def risky_forward(self, x):
    x = self.linear1(x)
    x.relu_()  # 学習時にこれをするのは危険
    x = self.linear2(x)
    return x

このコードの問題点は、バックプロパゲーション時に必要な中間結果が失われてしまうことです。代わりに、以下のような実装が安全です

def safe_forward(self, x):
    x = self.linear1(x)
    x = F.relu(x)  # 学習時は通常のReLUを使用
    x = self.linear2(x)
    return x

このように、インプレース操作は使用するタイミングがきわめて重要です。

このcustom_sgd_updateの実装は、PyTorchの公式オプティマイザーと同様の方針で、勾配計算が完了した後の重み更新フェーズで安全にインプレース操作を使用している例となっています。このように、適切なタイミングでインプレース操作を使用することで、メモリ効率を維持しながら安全な学習処理を実現することができます。

さて、くどくどと説明しましたが、読者様の理解が深まれば幸いです^o

まとめ：こんな感じで使っていきましょう！

これまでの内容を踏まえて、インプレース操作の実践的な活用方針についてまとめてみましょう。

まず、推論時の活用について見ていきましょう。

本番環境での推論では、バックプロパゲーションを考慮する必要がないため、積極的にインプレース操作を活用できます。

例えば、画像認識モデルの推論時には以下のようなコードが有効です

def inference(self, x):
    x = self.conv(x)
    x = self.bn(x)
    x.relu_()  # 推論時は安心してインプレース操作が使える
    x = self.pool(x)
    x.relu_()  # 複数回使っても問題なし
    return x

また、データの前処理でも、推論時には積極的に活用できます。例えば画像の正規化処理では、以下のように活用することができます

def preprocess_image(image_tensor):
    image_tensor.sub_(mean).div_(std)  # インプレースで正規化
    image_tensor.clamp_(min=0)  # 値の範囲も直接制限
    return image_tensor

一方、学習時は使用するタイミングを慎重に選ぶ必要があります。

モデルの重み更新時には、以下のようにインプレース操作を安全に使用できます。

def update_model(self):
    with torch.no_grad():
        for param in self.parameters():
            if param.grad is not None:
                param.sub_(learning_rate * param.grad)  # 重み更新時はOK

本番環境での実装では、これらの知識を組み合わせることで効率的なシステムを構築できます。

例えば、本番用のモデルクラスを以下のように実装すると良いですね

class ProductionModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.backbone = create_backbone()
        self.is_training = False  # 学習/推論モードのフラグ
    
    def forward(self, x):
        # 推論時はメモリ効率重視で実装
        if not self.is_training:
            x = self.preprocess_inplace(x)
            features = self.extract_features_inplace(x)
            return self.postprocess_inplace(features)
        
        # 学習時は通常の演算を使用
        x = self.preprocess_standard(x)
        features = self.extract_features_standard(x)
        return self.postprocess_standard(features)

インプレース操作を適切に活用することで、特に本番環境での推論において大きなメモリ効率の向上が期待できます。

ResNetのような深いネットワークでは、活性化関数やバッチ正規化でインプレース操作を使用することで、中間層のメモリ使用量を大幅に削減できます。

一方で学習時には、計算グラフの構築が完了した後の重み更新フェーズでのみインプレース操作を使用することで、安全性と効率性の両立が可能です。

このように、用途とタイミングを適切に見極めることで、インプレース操作は強力なメモリ最適化のツールとなりそうです。

GPUはメモリ量に応じて指数関数的に高額になる傾向にあり、ディープラーニングを活用した学習モデルの大規模化が進む中、こうした最適化テクニックの重要性・有用性は今後さらに高まりそうですので、インプレース操作をマスターして活用していきましょう！

Appendix

良く使うインプレース操作

基本的な計算系

x.add_(value) # 足し算
x.sub_(value) # 引き算
x.mul_(value) # 掛け算
x.div_(value) # 割り算

ニューラルネット系

x.relu_() # ReLU
x.sigmoid_() # Sigmoid
x.tanh_() # Tanh

その他

PyTorchにおけるインプレース演算子には、基本的な算術演算から高度な数学的操作まで、多様な種類が用意されています。

累乗を計算するpow_や、負の値に変換するneg_も利用可能です。

行列演算では、行列積を計算するmatmul_、テンソルの転置を行うtranspose_、次元の並び替えを行うpermute_などが提供されています。

要素ごとの操作としては、最小値でクリッピングするclamp_min_、最大値でクリッピングするclamp_max_、値の範囲を制限するclamp_があります。また、要素ごとの最小値を取るminimum_、最大値を取るmaximum_も用意されています。

データの初期化や変更に関する操作として、全ての要素をゼロにするzero_、指定した値で埋めるfill_、他のテンソルからデータをコピーするcopy_があります。また、ランダムな値で埋めるrandom_や、正規分布に従う乱数で埋めるnormal_も提供されています。

これらの演算子は、通常の演算子と同じように使用できますが、末尾にアンダースコアがついているという命名規則によって、インプレース操作であることを明示的に示しています。例えば、通常の加算がtensor + valueやtensor.add(value)と書くのに対し、インプレース加算はtensor.add_(value)のように書きます。これにより、コードを読む際にもメモリを直接書き換える操作であることが一目で分かるようになっています。

参考文献

PyTorch公式ドキュメント
"Deep Learning with PyTorch" ( Eli Stevens)
"Programming PyTorch for Deep Learning"( Ian Pointer)