Hacktoberfest 2026：メンテナが10月に向けて印を付けた、オープンで初心者向けの issue。 Hacktoberfest の issue を見る

[torchlib] aten_lstm / aten_gru: runtime shape computation between layers makes multi-layer RNNs ~10% slower

クローズ 初心者向け
  • 何か問題がありますか？
#3,096 コメント 0 件 リアクション 1 件 担当者 2 名 GitHub で見る

メンテナーはふだん 3 日以内に返信

@justinchuby がすでに取り組んでいます。

2026年10月20日 から。

  • #3111 @copilot-swe-agent による — マージ済み

評価

難易度
2/5
見積もり時間
1〜3時間
初心者へのやさしさ
68/100
issue の種類
リファクタリング
明瞭さ
明確に書かれている
活発さ
活発
技術スタック
python, pytorch
領域
machine-learning, performance

調査の方向性

onnxscript/function_libs/torch_lib/ops/core.py の aten_lstm と aten_gru から始めてください。ここでは各層の出力が次の層の入力に reshape されています。bidirectional が False の場合は方向軸を squeeze し、それ以外の場合は定数 [0, 0, 2 * hidden_size] による reshape を行うヘルパーを追加してください。双方向のケースを含め、LSTM と GRU の torch_lib テストを実行し、出力が変わらないこと、および実行時の Shape/Slice ノードが消えていることを確認してください。

索引モデルが issue の本文から書いたものです。

説明

Summary

aten_lstm and aten_gru (onnxscript/function_libs/torch_lib/ops/core.py, introduced in #2674) convert each layer's ONNX output
Y: [seq_length, num_directions, batch_size, hidden_size] into the next layer input
[seq_length, batch_size, num_directions * hidden_size] with:

Transpose(perm=[0, 2, 1, 3]) -> Shape -> 4x Slice -> Mul -> Reshape(-1) -> Concat -> Reshape

That is 10 nodes per layer, including a target shape computed at runtime. Two things about the shape are known when the graph is traced:

  • bidirectional is a Python bool. When it is False, num_directions == 1, so the whole sequence is just Squeeze(Y, axes=[1]), which is what the TorchScript exporter emits.
  • hidden_size is a static int (it is already passed as the LSTM/GRU hidden_size attribute). When bidirectional is True, the target shape can be the constant [0, 0, 2 * hidden_size], where 0 copies the input dimension.

With a dynamic batch axis, ONNX Runtime cannot constant-fold the shape computation, so every layer runs these extra tiny kernels.
In our case, it made the dynamo export slower than the TorchScript export of the same model (can be reproduced with dynamo=False).

Repro

Versions: torch 2.13.0+cpu, onnxscript main @ 9ff5a7c (same code in 0.7.1 / 0.7.2), onnxruntime 1.31.0, CPU, 1 thread.

import collections, sys, time
import numpy as np, onnx, onnxruntime as ort, torch


class StepLSTM(torch.nn.Module):
    def __init__(self, input_size=256, hidden_size=256, num_layers=4):
        super().__init__()
        self.lstm = torch.nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)

    def forward(self, x, h0, c0):
        return self.lstm(x, (h0, c0))


torch.manual_seed(0)
model = StepLSTM().eval()
batch = 5
inputs = (torch.randn(batch, 1, 256), torch.randn(4, batch, 256), torch.randn(4, batch, 256))
batch_dim = torch.export.Dim("batch")
path = sys.argv[1]
torch.onnx.export(
    model, inputs, path, dynamo=True,
    input_names=["x", "h0", "c0"],
    dynamic_shapes=({0: batch_dim}, {1: batch_dim}, {1: batch_dim}),
)

counts = collections.Counter(node.op_type for node in onnx.load(path).graph.node)
print(f"nodes={sum(counts.values())} {dict(sorted(counts.items()))}")

options = ort.SessionOptions()
options.intra_op_num_threads = 1
options.inter_op_num_threads = 1
session = ort.InferenceSession(path, options, providers=["CPUExecutionProvider"])
feeds = {name: tensor.numpy() for name, tensor in zip(["x", "h0", "c0"], inputs)}
for _ in range(500):
    session.run(None, feeds)
timings = []
for _ in range(5):
    start = time.perf_counter()
    for _ in range(2000):
        session.run(None, feeds)
    timings.append((time.perf_counter() - start) / 2000 * 1e6)
print(f"ort_latency={min(timings):.1f} us/step (best of 5)")
Results
nodes Reshape Shape Slice ORT latency (us/step)
main 104 8 4 56 991 - 1014
with the proposed change 68 0 0 40 908 - 911

Outputs are identical: max abs diff vs PyTorch is 8.3e-07 in both cases. That is about 9-10% faster for this model

Proposed fix

Factor the conversion into a helper used by both aten_lstm and aten_gru:

def _rnn_output_to_layer_output(Y: TFloat, bidirectional: bool, hidden_size: int) -> TFloat:
    """Y: [seq_length, num_directions, batch_size, hidden_size]
    Returns: [seq_length, batch_size, num_directions * hidden_size]
    """
    if not bidirectional:
        # num_directions == 1: dropping the axis is enough
        return op.Squeeze(Y, [1])
    # [seq_length, batch_size, num_directions, hidden_size]
    Y = op.Transpose(Y, perm=[0, 2, 1, 3])
    # 0 copies the input dimension (seq_length, batch_size): no runtime shape computation needed
    return op.Reshape(Y, op.Constant(value_ints=[0, 0, 2 * hidden_size]))

onnxscript-rnn-squeeze.patch

Image Image
主要言語
Python
スター
456
フォーク
148
平均マージ
3日 19時間
マージ済み PR（30日）
20

環境構築

はじめの一歩

  1. issue を最後まで読み、次にプロジェクトのコントリビューションガイドを読みます。
  2. 着手することを issue にコメントします — 二人が同じ作業をするのを防げます。
  3. リポジトリをフォークし、ブランチを切って変更します。
  4. issue 番号を参照したプルリクエストを送ります。

microsoft/onnxscript のほかの issue

microsoft/onnxscript の issue をすべて見る

似ている issue

Python の issue をもっと見る

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。