Fix autograd in-place error: move memory-bank enqueue after backward
Hard mining (c307269) uses `queue.t()` as a view directly in the
similarity matmul, so the autograd graph now holds a reference to the
memory-bank buffer's storage. The previous code enqueued fresh gallery
embeddings BEFORE `backward()`, which mutates that same buffer in place
and triggers:
RuntimeError: one of the variables needed for gradient computation has
been modified by an inplace operation: [torch.cuda.FloatTensor [1024, 8]]
is at version 2; expected version 1 instead.
The older InfoNCE path called `torch.cat([emb_b, queue], 0)`, which
materialises a fresh contiguous tensor and severs the alias — so the
same enqueue order worked. After the mining refactor, we need the
buffer to stay stable until backward completes.
Moving the enqueue to after backward is semantically identical: the
queue state observed by the next training step is the same either way
(FIFO, not involved in the just-completed backward).
`smoke_train.py` already had the correct order and didn't catch the
regression.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -839,15 +839,20 @@ def train(cfg: TrainConfigGTAUAV) -> None:
|
||||
)
|
||||
loss_dict = loss_fn(**loss_kwargs)
|
||||
|
||||
# Enqueue current gallery embeddings (detached).
|
||||
if neg_bank is not None:
|
||||
neg_bank.enqueue(embeddings["gallery"].detach())
|
||||
|
||||
# Scale loss by accumulation steps so gradients average correctly.
|
||||
raw_loss = float(loss_dict["total"].item()) # save before backward
|
||||
total_loss = loss_dict["total"] / accum
|
||||
scaler.scale(total_loss).backward()
|
||||
|
||||
# Enqueue current gallery AFTER backward. The queue buffer is aliased
|
||||
# into the autograd graph through `queue_neg` (a view returned by
|
||||
# `NegativeMemoryBank.get_queue`), so modifying it before backward
|
||||
# triggers "variable needed for gradient computation has been modified
|
||||
# by an inplace operation". Enqueueing here is semantically identical
|
||||
# — the next step's queue state is the same either way.
|
||||
if neg_bank is not None:
|
||||
neg_bank.enqueue(embeddings["gallery"].detach())
|
||||
|
||||
# Optimizer step only after accumulating `accum` micro-batches.
|
||||
is_accum_step = (n_batches + 1) % accum == 0 or (n_batches + 1) == len(train_loader)
|
||||
if is_accum_step:
|
||||
|
||||
Reference in New Issue
Block a user