チートシート
テンソルの形状とブロードキャスト
形状の対応表で見るブロードキャスト規則、注意機構のためのeinsumパターン、reshapeとpermuteの違い、バッチ次元、そして注意マスク、畳み込み、損失関数に潜む形状のバグを扱います。
このシートは、RuntimeError: shape mismatchを睨んだことがある人、あるいはもっと悪いことに、エラーになるべきところでテンソルがブロードキャストされたせいで、学習は順調なのに結果が静かに間違っているモデルを前にしたことがある人のためのものです。規則はNumPy、PyTorch、JAXで共通です。ブロックはCPU上で1秒以内に動きます。
ブロードキャスト規則
2つの形状は右から比較されます。各次元の組は等しいか、どちらかが1（または存在しない）でなければなりません。結果は各位置で大きいほうのサイズを取ります。規則はそれ以外にありません。
| 左 | 右 | 右揃えにすると | 結果 | 理由 |
|---|---|---|---|---|
(8, 1, 4) |
(5, 1) |
(8, 1, 4)と(1, 5, 1) |
(8, 5, 4) |
欠けた次元は1になり、1が引き伸ばされる |
(3, 4) |
(4,) |
(3, 4)と(1, 4) |
(3, 4) |
行ベクトルがすべての行に加算される |
(3, 4) |
(3,) |
(3, 4)と(1, 3) |
エラー | 4と3は異なり、どちらも1ではない。(3, 1)を使う |
(3, 4) |
(3, 1) |
同じ | (3, 4) |
列ベクトルがすべての列に加算される |
(32, 10) |
(32,) |
(32, 10)と(1, 32) |
バッチが10でない限りエラー | サンプルごとの重みにまつわる典型的なバグ。w[:, None]を使う |
(2, 1, 5, 5) |
(1, 4, 5, 5) |
同じ | (2, 4, 5, 5) |
バッチとヘッドの次元にまたがるマスク |
(N, C) |
(C,) |
(N, C)と(1, C) |
(N, C) |
バイアスやクラスごとのスケール |
import numpy as np
a = np.ones((8, 1, 4)) # rank 3
b = np.ones((5, 1)) # rank 2, right-aligned to (1, 5, 1)
c = a + b # (8, 1, 4) with (1, 5, 1) -> (8, 5, 4)
print("a", a.shape, "+ b", b.shape, "->", c.shape)
x = np.arange(12.0).reshape(3, 4)
mean_per_row = x.mean(axis=1, keepdims=True) # (3, 1), broadcasts back over columns
print("centred rows:", (x - mean_per_row).shape)
try:
np.ones((3, 4)) + np.ones((4, 3))
except ValueError as err:
print("error:", err)a (8, 1, 4) + b (5, 1) -> (8, 5, 4)
centred rows: (3, 4)
error: operands could not be broadcast together with shapes (3,4) (4,3)keepdims=Trueは、ブロードキャストのバグの大半を避ける習慣です。リダクションがランクを保つので、Noneやunsqueezeなしで結果が入力と揃います。危険なのは形状がたまたま揃ってしまう場合です。たとえば(32,)のサンプルごとの重みを(32, 32)のペアワイズスコアに掛けると、間違った軸でブロードキャストされたまま学習が進みます。
バッチ次元と覚えておくべき慣習
| データ | PyTorchのレイアウト | 補足 |
|---|---|---|
| 表形式 | (N, features) |
|
| トークン列 | IDは(N, T)、埋め込みは(N, T, D) |
nn.LSTMとnn.MultiheadAttentionではbatch_first=True、そうでなければ(T, N, D) |
| 画像 | (N, C, H, W) |
PILとほとんどの画像ライブラリは(H, W, C)を返す。TensorFlowの既定は(N, H, W, C) |
| 音声 | (N, channels, samples)、スペクトログラム後は(N, T, mel_bins) |
|
| 注意スコア | (N, heads, T_q, T_k) |
マスクは(T_q, T_k)または(N, 1, T_q, T_k) |
| 分類のロジット | (N, classes) |
ラベルはone-hotではなくint64の(N,) |
| トークン予測のロジット | (N, T, vocab) |
cross_entropyには(N*T, vocab)と(N*T,)に平坦化するか、(N, vocab, T)を渡す |
einsumパターン
einsumは各軸に文字を付け、総和を取るべき軸には同じ文字を繰り返し、出力の順序を列挙します。transpose、unsqueeze、bmmの連鎖を、形状の宣言として読み返せる1行に置き換えます。
| パターン | 読み方 | 等価な式 |
|---|---|---|
"ij,jk->ik" |
行列積 | a @ b |
"bij,bjk->bik" |
バッチ行列積 | torch.bmm(a, b) |
"btd,de->bte" |
全トークンへの線形層 | x @ w |
"bhqd,bhkd->bhqk" |
全ヘッド同時の注意スコア | q @ k.transpose(-2, -1) |
"bhqk,bhkd->bhqd" |
注意重み付きの値の和 | attn @ v |
"bi,bj->bij" |
バッチ外積 | a[:, :, None] * b[:, None, :] |
"ii->" |
トレース | torch.trace(a) |
"bij->bji" |
バッチ転置 | a.transpose(1, 2) |
"bd,bd->b" |
行ごとの内積 | (a * b).sum(-1) |
import torch
torch.manual_seed(0)
B, T, D, H = 2, 5, 8, 2 # batch, tokens, model dim, heads
x = torch.randn(B, T, D)
w = torch.randn(D, D)
# matrix multiply: (B,T,D) x (D,D) -> (B,T,D)
y = torch.einsum("btd,de->bte", x, w)
# split heads: (B,T,D) -> (B,H,T,D/H)
q = y.view(B, T, H, D // H).transpose(1, 2)
k, v = q.clone(), q.clone()
# attention scores per head: (B,H,T,d) x (B,H,T,d) -> (B,H,T,T)
scores = torch.einsum("bhqd,bhkd->bhqk", q, k) / (D // H) ** 0.5
attn = scores.softmax(dim=-1)
# weighted sum of values: (B,H,T,T) x (B,H,T,d) -> (B,H,T,d)
out = torch.einsum("bhqk,bhkd->bhqd", attn, v)
merged = out.transpose(1, 2).reshape(B, T, D)
# the same with matmul, to check
scores2 = (q @ k.transpose(-2, -1)) / (D // H) ** 0.5
print("scores", scores.shape, "merged", merged.shape, "einsum == matmul:", torch.allclose(scores, scores2))
# batched outer product and trace are one line each
outer = torch.einsum("bi,bj->bij", torch.randn(3, 4), torch.randn(3, 6))
trace = torch.einsum("ii->", torch.eye(4))
print("outer", outer.shape, "trace", trace.item())scores torch.Size([2, 2, 5, 5]) merged torch.Size([2, 5, 8]) einsum == matmul: True
outer torch.Size([3, 4, 6]) trace 4.0最後のtransposeの後の.reshapeに注目してください。ヘッドは論理的な順序では隣り合わせに移動しましたが、メモリ上ではそうなっていません。それが次の節の話です。
reshapeとpermute
reshapeとviewはメモリ上の順序を保ったまま、次元への分け方を変えます。permuteとtransposeはどの軸がどれかを入れ替えます。両者は互換ではなく、混同すると形状は正しいのに中身がかき混ぜられたテンソルができ、どんなエラーも捕まえてくれません。
import torch
t = torch.arange(6).reshape(2, 3)
print(t)
print("reshape(3, 2):\n", t.reshape(3, 2)) # reads elements in memory order, row 0 then row 1
print("permute(1, 0):\n", t.permute(1, 0)) # transposes: element [i, j] moves to [j, i]
p = t.permute(1, 0)
print("permuted is contiguous:", p.is_contiguous())
try:
p.view(-1)
except RuntimeError as err:
print("view fails:", str(err).split(".")[0])
print("contiguous().view works:", p.contiguous().view(-1).tolist())
print("reshape copies when needed:", p.reshape(-1).tolist())reshape(3, 2)は行[0, 1], [2, 3], [4, 5]を返します。permute(1, 0)は行[0, 3], [1, 4], [2, 5]を返します。同じ形状で、異なるテンソルです。permuteの後はテンソルが非連続になり、viewは拒否し、reshapeは黙ってコピーします。
| やりたいこと | 使うもの | 使わないもの |
|---|---|---|
画像(H, W, C)を(C, H, W)に |
img.permute(2, 0, 1) |
img.reshape(C, H, W)（ピクセルがかき混ざる） |
ヘッドの結合(B, H, T, d)を(B, T, H*d)に |
x.transpose(1, 2).reshape(B, T, H * d) |
x.reshape(B, T, H * d)（トークンとヘッドが混ざる） |
| 線形層のために画像バッチを平坦化 | x.flatten(1) |
x.view(-1)（バッチ軸が失われる） |
| バッチ軸を追加 | x.unsqueeze(0)またはx[None] |
x.reshape(1, -1)（平坦化もしてしまう） |
| 系列とバッチを入れ替え | x.transpose(0, 1) |
x.reshape(T, N, D) |
注意機構と畳み込みの形状バグ
import torch
from torch import nn
torch.manual_seed(0)
# Convolutions: PyTorch wants (N, C, H, W). An image loaded as (H, W, C) has to be permuted.
img_hwc = torch.rand(32, 32, 3)
x = img_hwc.permute(2, 0, 1).unsqueeze(0) # (1, 3, 32, 32)
conv = nn.Conv2d(3, 8, kernel_size=3, stride=2, padding=1)
out = conv(x)
assert out.shape == torch.Size([1, 8, 16, 16]), out.shape
print("conv out", tuple(out.shape)) # floor((32 + 2 - 3) / 2) + 1 = 16
# Attention: a mask of the wrong rank broadcasts silently. Pin the shapes.
B, H, T = 2, 4, 6
scores = torch.randn(B, H, T, T)
causal = torch.tril(torch.ones(T, T, dtype=torch.bool)) # (T, T)
assert causal.shape == (T, T)
masked = scores.masked_fill(~causal, float("-inf")) # broadcasts over B and H
probs = masked.softmax(dim=-1)
assert torch.allclose(probs.sum(-1), torch.ones(B, H, T))
print("row 0 attends to", int((probs[0, 0, 0] > 0).sum()), "token(s); row 5 attends to", int((probs[0, 0, 5] > 0).sum()))
# Loss functions have shape contracts too
logits = torch.randn(8, 10) # (N, classes)
labels = torch.randint(0, 10, (8,)) # (N,) integer class ids, not one-hot
print("cross-entropy", round(nn.functional.cross_entropy(logits, labels).item(), 3))conv out (1, 8, 16, 16)
row 0 attends to 1 token(s); row 5 attends to 6
cross-entropy 2.883畳み込みの出力サイズは、空間軸ごとにfloor((H + 2 * padding - dilation * (kernel - 1) - 1) / stride) + 1です。層の隣にコメントとして書いておきましょう。assertは実行時のコストがほぼゼロで、かき混ざった結果をトレースバックに変えてくれます。
| バグ | 症状 | 確認方法 |
|---|---|---|
(N, 1, 1, T)が必要な場所に(N, T)のパディングマスクを適用 |
エラー、または間違った軸をマスクする | mask[:, None, None, :]にしてランク4をassertする |
因果マスクをtrilではなくtriuで作成 |
モデルが未来を参照し、損失が不自然に低い | 左下と右上の角でassert mask[0, -1] and not mask[-1, 0] |
dim=-2でのsoftmax |
行ではなく列の和が1になる | assert torch.allclose(probs.sum(-1), 1) |
nn.MultiheadAttentionにbatch_first=Trueなしで(N, T, D)を渡す |
バッチと時間が入れ替わり、学習がうまくいかない | batch_first=Trueを渡し、出力形状が入力形状と等しいことをassertする |
channels-lastの画像をConv2dに渡す |
Cが3でなければエラー、Hがたまたま3ならゴミが出る | assert x.shape[1] == in_channels |
cross_entropyにラベル形状(N, 1) |
エラー、またはBCEWithLogitsLossでは(N, 1)対(N,)が(N, N)にブロードキャストされる警告 |
labels.squeeze(-1)し、labels.shape == logits.shape[:1]をassertする |
サイズ1のバッチが.squeeze()で消える |
最後のバッチまでは動く | 軸を明示したsqueeze(dim) |
torch.Sizeによるアサーション
tensor.shapeはtorch.Sizeで、tupleのサブクラスなので普通のタプルと比較できます。モデルの各段階の後に置く1行のassertは、どんなデバッガより安上がりです。コードに残しておきましょう。契約を文書化し、コストはマイクロ秒です。変動する形状については、ランクと固定次元をassertします。たとえばassert x.ndim == 3 and x.shape[-1] == d_modelです。より長く使う契約には、関数シグネチャのdocstringで次元に名前を付けるか、形状を文書化し実行時に検査もできるjaxtypingアノテーションを使います。