ailia-models
ailia-models copied to clipboard
axinc-ai
Improve performance of RAFT
@mucunwuxian RAFTのコードを見ていまして、RAFTはITERS回数分だけ、Flowを更新すると考えています。ここで、cnetはITERごとにrunを実行しているのですが、入力がimage1にしか依存しないような気がしていまして、cnet.runをITERSの外に出して、一度だけ呼び出すようにしても問題はないものでしょうか。
flow_predictions = []
for itr in range(ITERS):
corr = corr_fn(coords1) # index correlation volume
flow = coords1 - coords0
logger.info('[iter:%d] Start predicting optical flow...' % (itr+1))
cmap = cnet.run(image1)[0]
net, up_mask, delta_flow = update_block.run([net, inp, corr, flow])
# F(t+1) = F(t) + \Delta(t)
coords1 = coords1 + delta_flow
# upsample predictions
if (args.model == 'small'):
flow_up = upflow8(coords1 - coords0)
else:
flow_up = upsample_flow(coords1 - coords0, up_mask)
flow_predictions.append(flow_up)
M1 Maxでの処理時間。
itr processing time 968 ms
+ corr_fn processing time 316 ms
+ cnet processing time 275 ms
+ update_block processing time 313 ms
+ upsample processing time 64 ms
-m smallの場合の処理時間。
itr processing time 298 ms
+ corr_fn processing time 183 ms
+ cnet processing time 45 ms
+ update_block processing time 66 ms
+ upsample processing time 4 ms
@kyakuno こちら、返信遅れて申し訳ありません…！🙇💦
ご確認、及び、ご指摘頂き、ありがとうございます。 私の方でも確認をさせて頂こうと思います。
@kyakuno こちら、解決までに時間がかかってしまい、申し訳ありません…。🙇💦
客野さんの仰る通り、loop内のcmapの呼び出しにつきましては、完全にコーディングミスでした。 すぐ上部にある [BENCHMARK mode] のコードを踏襲すべくコピーした後、cmapの計算コードを削除せずに残してしまっていた次第でした。
尚、踏襲元のコードにつきましても、cmapの計算は1度きりとなっております。
修正後は、以下のようなコードとなります。
# predict optical flow
flow_predictions = []
for itr in range(ITERS):
corr = corr_fn(coords1) # index correlation volume
flow = coords1 - coords0
logger.info('[iter:%d] Start predicting optical flow...' % (itr+1))
if args.benchmark:
logger.info('BENCHMARK mode')
for i in range(args.benchmark_count):
start = int(round(time.time() * 1000))
net, up_mask, delta_flow = update_block.run([net, inp, corr, flow])
end = int(round(time.time() * 1000))
logger.info(f'\tailia processing time {end - start} ms')
else:
net, up_mask, delta_flow = update_block.run([net, inp, corr, flow])
# F(t+1) = F(t) + \Delta(t)
coords1 = coords1 + delta_flow
# upsample predictions
if (args.model == 'small'):
flow_up = upflow8(coords1 - coords0)
else:
flow_up = upsample_flow(coords1 - coords0, up_mask)
flow_predictions.append(flow_up)
以上。
現在、上記修正前後におきまして、ある1つのサンプル画像からのアウトプットが同値であることまでは、確認が取れています。 この後、もう少し精緻に確認をさせて頂き、問題無さそうでしたらPRを出させて頂こうと思います。
尚、計算時間ですが、修正前後で10%程高速化がされた次第です。
- 通常モデル：
- 修正前：30.45347 (sec)
- 修正後：26.48193 (sec)
- smallモデル：
- 修正前：8.72459 (sec)
- 修正後：8.17172 (sec)
他にも、高速化の余地が無いかについても、少しだけ検討させて頂こうと思います。