私はいくつかのサンプルを転送するだけで、計算グラフはいつ解放されますか？

Question

私はバッチ内の各サンプルについてフォワードを行い、サンプルのモデル出力上のいくつかの条件に基づいて一部のサンプルについてのみ損失を累積するユースケースを持っています。ここにコードを例示します。

for batch_idx, (data, target) in enumerate(train_loader): 
    optimizer.zero_grad() 
    total_loss = 0 

    loss_count_local = 0 
    for i in range(len(target)): 
     im = Variable(data[i].unsqueeze(0).cuda()) 
     y = Variable(torch.FloatTensor([target[i]]).cuda()) 

     out = model(im) 

     # if out satisfy some condtion, we will calculate loss 
     # for this sample, else proceed to next sample 
     if some_condition(out): 
      loss = criterion(out, y) 
     else: 
      continue 

     total_loss += loss 
     loss_count_local += 1 

     if loss_count_local == 32 or i == (len(target)-1): 
      total_loss /= loss_count_local 
      total_loss.backward() 
      total_loss = 0 
      loss_count_local = 0 

    optimizer.step() 

私の質問は、すべてのサンプルについては転送しますが、いくつかのサンプルについては後で行います。損失に寄与しないサンプルのグラフはいつ解放されますか？これらのグラフは、forループが終了した後、または次のサンプルを転送した直後に解放されますか？私はここで少し混乱しています。

total_lossに寄与するサンプルについては、total_loss.backward()の直後にグラフが解放されます。そうですか？

Answer 1

がPyTorchがメモリを解放する方法の一般的な議論を始めるのをしてみましょう：

まず、我々はPyTorchはPythonオブジェクトの属性に格納され、暗黙的に宣言したグラフを使用していることを強調すべきです。 （覚えているのはPythonだから、すべてがオブジェクトだからです）。具体的には、torch.autograd.Variableは.grad_fn属性を持っています。この属性の型は、どのような種類の計算ノード（例えば、加算）と、そのノードへの入力を定義する。

これはPytorchが標準的なpythonガベージコレクタを使用するだけで（かなり積極的に）メモリを解放するので重要です。このコンテキストでは、これは、（暗黙的に宣言された）計算グラフが、現在のスコープでそれらを保持するオブジェクトへの参照が存在する限り、生き続けることを意味します！

これは、たとえばサンプルs_1 ... s_kでいくつかの種類のバッチ処理を行い、それぞれの損失を計算して最後に損失を加えます。その累積損失は、計算された各計算ノードへの参照を保持する個々の損失を参照しますそれ。

コードに適用される質問は、Python（または、より具体的にはガベージコレクタ）がPytorchについてよりも参照を処理する方法の詳細です。 1つのオブジェクト（total_loss）に損失が蓄積されるため、ポインタを有効にしておくことで、外側のループでそのオブジェクトを再初期化するまでメモリを解放しません。

これは、往路で作成した計算グラフ（out = model(im)）は、outオブジェクトとそれ以降の計算によって参照されることを意味します。したがって、損失を計算して合計すると、outの参照が生き残り、計算グラフに残ります。ただし、それを使用しない場合は、ガベージコレクタはoutとその計算グラフを再帰的に収集する必要があります。