d2l-ai/d2l-zh · error · AssertionError

X.shape[0] == y.shape[0]

Error message

X.shape[0] == y.shape[0]

What it means

AssertionError from d2l.mxnet.split_batch (sec_multi_gpu): features X and labels y must have the same batch dimension before being scattered with gluon.utils.split_and_load across GPUs. It is a cheap pre-flight check; failure means the data pairing is broken before multi-GPU training starts.

Source

Thrown at d2l/mxnet.py:1368

class Benchmark:
    """用于测量运行时间"""
    def __init__(self, description='Done'):
        """Defined in :numref:`sec_hybridize`"""
        self.description = description

    def __enter__(self):
        self.timer = d2l.Timer()
        return self

    def __exit__(self, *args):
        print(f'{self.description}: {self.timer.stop():.4f} sec')

def split_batch(X, y, devices):
    """将X和y拆分到多个设备上

    Defined in :numref:`sec_multi_gpu`"""
    assert X.shape[0] == y.shape[0]
    return (gluon.utils.split_and_load(X, devices),
            gluon.utils.split_and_load(y, devices))

def resnet18(num_classes):
    """稍加修改的ResNet-18模型

    Defined in :numref:`sec_multi_gpu_concise`"""
    def resnet_block(num_channels, num_residuals, first_block=False):
        blk = nn.Sequential()
        for i in range(num_residuals):
            if i == 0 and not first_block:
                blk.add(d2l.Residual(
                    num_channels, use_1x1conv=True, strides=2))
            else:
                blk.add(d2l.Residual(num_channels))
        return blk

    net = nn.Sequential()

View on GitHub (pinned to e6b18ccea7)

Solutions

  1. Align the slices: use the same index/length for X and y (X[:n], y[:n])
  2. Take X, y as a pair from the same DataLoader batch instead of assembling them from separate sources
  3. Add a guard `if X.shape[0] != y.shape[0]: raise ValueError(...)` with context (shapes, batch idx) before calling split_batch for clearer diagnostics
  4. Verify no stray reshape dropped the batch axis (e.g. reshape(y, -1) versus reshape(y, (batch, 1)))

Example fix

# before
X, y = next(iter(train_iter))
X = X[:64]  # trimmed features only
split_batch(X, y, devices)  # y has 256 rows -> AssertionError
# after
X, y = X[:64], y[:64]
split_batch(X, y, devices)  # both 64 rows, scatter ok
Defensive patterns

Strategy: validation

Validate before calling

assert X.shape[0] == y.shape[0], f'batch mismatch: X={X.shape}, y={y.shape}'
Xs, ys = split_batch(X, y, devices)

Type guard

def batch_aligned(X, y) -> bool:
    return getattr(X, 'shape', None) is not None and X.shape[0] == y.shape[0]

Prevention

When it happens

Trigger: split_batch(X[:32], y[:16], devices) via mismatched slicing; a DataLoader returning y=None or a padded y of different length; feeding X from one iterator and y from another (async prefetch bugs); passing tensors whose batch dims diverged after a transpose/reshape mistake.

Common situations: Hand-rolled multi-GPU training loops adapted from the book; batch trimming like X = X[:batch_size] applied to only one of X/y; mixed torch tensors and mxnet ndarrays where .shape semantics surprise the user; shuffling X and y separately so rows no longer correspond.

Related errors


AI-assisted analysis of d2l-ai/d2l-zh@e6b18ccea7 (2026-08-14). Data as JSON: /api/errors/6b49ed52bd4b5812. Report an issue: GitHub.