PyTorch の ResNet-18 モデルにいくつかの変更を加えて、各 ResNet ブロックの最後にある ResNet 中間層出力を入力として取り込み、推論中にいくつかの補助予測を行う別の補助トレーニング モデルの実行を呼び出そうとしています。フェーズ。
GPU でのパイプライン実行全体のエンドツーエンドのレイテンシを削減するために、次の ResNet ブロックの計算と並行してブロックの計算の後に補助計算を実行できるようにしたいと考えています。
機能の観点からは正しく動作する基本コードがありますが、補助モデルの実行は ResNet ブロックの計算と連続しています。これを2つの方法で確認しました-
print ステートメントを追加し、実行順序を検証します。
元の ResNet モデルの実行時間 (時間 t1 とします) と補助モデル (時間 t2 とします) を計測します。私の実行時間は現在 t1+t2 です。
元の ResNet ブロック コード (ResNet-18 で実験しているため、これは BasicBlock です)。コード全体はこちらから入手できます
class BasicBlock(nn.module):
...
def forward(self, x):
residual = x
out = self.conv1(x)
out = self.bn1(out)
out = self.relu(out)
out = self.conv2(out)
out = self.bn2(out)
if self.downsample is not None:
residual = self.downsample(x)
out += residual
out = self.relu(out)
return out
これはシリアル方式で動作する私の修正です -
def forward(self, x):
if len(x[0]) == self.auxiliary_prediction_size: # Got an Auxiliary prediction earlier
return x
# Do usual block computation
residual = x
out = self.conv1(x)
out = self.bn1(out)
out = self.relu(out)
out = self.conv2(out)
out = self.bn2(out)
if self.downsample is not None:
residual = self.downsample(x)
out += residual
out = self.relu(out)
# Try to make an auxiliary prediction
# First flatten the tensor (also assume for now that batch size is 1)
batchSize = x.shape[0]
intermediate_output = out.view(batchSize, -1)
# Place the flattened on GPU
device = torch.device("cuda:0")
input = intermediate_output.to(device)
# Make auxiliary prediction
auxiliary_input = out.float()
auxiliary_prediction = self.auxiliary_model(auxiliary_input)
if auxiliary_prediction meets some condition:
return auxiliary_prediction
# If no auxiliary prediction, then return intermediate output
return out
当然のことながら、上記のコードにより、補助モデルの実行と次のブロックの間でデータの依存関係が発生するため、物事が連続して発生します。私が最初に試した解決策は、このデータ依存関係を解消することでレイテンシが減少するかどうかを確認することでした。補助モデルの実行を許可し、条件が満たされた場合にauxiliary_predictionを返さないことでそうしようとしました（これは機能を壊しますが、この実験は純粋に動作を理解するためのものでした）. 本質的に、私がしたことは -
batchSize = x.shape[0]
intermediate_output = out.view(batchSize, -1)
# Place the flattened on GPU
device = torch.device("cuda:0")
input = intermediate_output.to(device)
# Make auxiliary prediction
auxiliary_input = out.float()
auxiliary_prediction = self.auxiliary_model(auxiliary_input)
if auxiliary_prediction meets some condition:
# Comment out return to break data dependency
#return auxiliary_prediction
# If no auxiliary prediction, then return intermediate output
return out
ただし、これは機能せず、さらに調査したところ、この Stack Overflow リンクで CUDA ストリームを見つけました。以下の方法で問題を解決するために、CUDAストリームのアイデアを取り入れてみました-
def forward(self, x):
if len(x[0]) == self.auxiliary_prediction_size: # Got an Auxiliary prediction earlier
return x
s1 = torch.cuda.Stream()
s2 = torch.cuda.Stream()
with torch.cuda.Stream(s1):
# Do usual block computation
residual = x
out = self.conv1(x)
out = self.bn1(out)
out = self.relu(out)
out = self.conv2(out)
out = self.bn2(out)
if self.downsample is not None:
residual = self.downsample(x)
out += residual
out = self.relu(out)
with torch.cuda.Stream(s2):
# Try to make an auxiliary prediction
# First flatten the tensor (also assume for now that batch size is 1)
out_detach = out.detach() # Detach from backprop flow and from computational graph dependency
batchSize = x.shape[0]
intermediate_output = out_detach.view(batchSize, -1)
# Place the flattened on GPU
device = torch.device("cuda:0")
input = intermediate_output.to(device)
# Make auxiliary prediction
auxiliary_input = out_detach.float()
auxiliary_prediction = self.auxiliary_model(auxiliary_input)
if auxiliary_prediction meets some condition:
return auxiliary_prediction
# If no auxiliary prediction, then return intermediate output
return out
ただし、Nvidia Visual Profiler からの出力は、すべての作業がまだデフォルト ストリームで実行され、シリアル化されていることを示しています。私が使用しているCUDAバージョンでCUDAストリームがサポートされていることを小さなCUDAプログラムで確認したことに注意してください。
私の質問 -
データの依存関係を壊しても、PyTorch が計算を並列にスケジュールしないのはなぜですか? これが PyTorch の動的計算グラフのポイントだと思いました。
CUDA ストリームを使用しても計算がデフォルト以外のストリームに委譲されないのはなぜですか?
補助モデルを ResNet ブロック計算と非同期/並列に実行する別の方法はありますか?