PyTorch はミニバッチ学習に自動的に対応する

PyTorch のチュートリアルのソースコードを見ていて，理解できない部分がありました．該当部分を抜粋して引用します．

class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28 * 28, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 10),
        )

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits

def train_loop(dataloader, model, loss_fn, optimizer):
    size = len(dataloader.dataset)
    # Set the model to training mode - important for batch normalization and dropout layers
    # Unnecessary in this situation but added for best practices
    model.train()
    for batch, (X, y) in enumerate(dataloader):
        # Compute prediction and loss
        pred = model(X) # <- ???
        loss = loss_fn(pred, y)

        # Backpropagation
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

        if batch % 100 == 0:
            loss, current = loss.item(), (batch + 1) * len(X)
            print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")

Tensor X は torch.Size([64, 1, 28, 28]) で，Model 内で最初に nn.Flatten() により torch.Size([64, 784]) となりますが，次に nn.Linear(28 * 28, 512) に渡されています．torch.Size([784]) となるべきで， shape が合致しないと思いました．しかし，実際は実行可能です．理解できませんでした．

Flatten — PyTorch 2.0 documentation

Linear — PyTorch 2.0 documentation

PyTorch の nn.Module では shape の最後の次元が適合していれば問題ないようで，それより上の次元についてはまとめて計算してくれるようです．

試しに，bias=False, weight=tensor([1.0]) の nn.Linear と nn.functional.relu() を使って，スカラを処理する単純なモデルを作成して動作を確認してみます．

class Model(nn.Module):
    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.layer = nn.Linear(1, 1, bias=False) # disable bias
        self.layer.weight.data = tensor([1], dtype=float32) # set weight to 1

    def forward(self, input: Tensor) -> Tensor:
        return nn.functional.relu(self.layer(input))

このモデルは，入力が負のスカラならば0を，正のスカラならば同じ値を返すだけです．

h = tensor([-1], dtype=float32)
a = model(h)

入力が torch.Size([1]) なので当然実行可能です．

# h
tensor([-1.])
# a
tensor(0., grad_fn=<ReluBackward0>)

次元を増やすとどうなるでしょうか．

i = tensor(range(-4, 5), dtype=float32).unsqueeze(-1)
b = model(i)

入力は torch.Size([9, 1]) ですが，実行できました．

# i
tensor([
        [-4.],
        [-3.],
        [-2.],
        [-1.],
        [ 0.],
        [ 1.],
        [ 2.],
        [ 3.],
        [ 4.]
])
# b
tensor([0., 0., 0., 0., 0., 1., 2., 3., 4.], grad_fn=<ReluBackward0>)

もちろん最後の次元の数が異なるとエラーになります．

j = tensor(range(-4, 5), dtype=float32)
c = model(j)

RuntimeError: inconsistent tensor size, expected tensor [9] and src [1] to have the same number of elements, but got 9 and 1 elements respectively

DataLoader はミニバッチが指定されていれば Tensor の上位の次元を一つ増やすという仕様になっており*1，上記の仕様と合わさると，自動的にミニバッチ学習ができる仕組みになっているようです．

*1:torch.utils.data — PyTorch 2.0 documentation より