实现梯度累积

Question

我正在使用下面的函数来训练我的模型一个时期。由于图像很大，我使用的是 8 的小批量大小以确保它们可以适合 GPU。为了模拟 256 的更大批量大小，我添加了梯度累积。然而，我注意到训练时间并没有减少，即使梯度累积步长很大。我不确定我的实施是否正确。有什么问题吗？

代码

def train_epoch(self):
    self.model.train()

    pbar = tqdm(unit="batch", file=sys.stdout, total=len(self.train_data))
    for batch_idx, data in enumerate(train_loader, start=1):
        inputs, labels = data

        inputs = inputs.to(self.device)
        labels = labels.to(self.device)

        with torch.autocast(self.device.type if self.device.type != 'mps' else 'cpu', enabled=self.amp):
            # Forward pass
            outputs = self.model(inputs)
            loss = F.cross_entropy(outputs, labels)

        # Backward pass
        self.grad_scaler.scale(loss).backward()

        # Accumulate gradients
        if batch_idx % self.gradient_accumulate_every == 0:
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.gradient_clipping)
            self.grad_scaler.step(self.optimizer)
            self.grad_scaler.update()
            self.optimizer.zero_grad(set_to_none=True)

Answer 1

Gradient Accumulation用于在大批量不适合内存时模仿大批量更新权重。 它不会加快你的计算速度。相比之下，它减慢了训练速度，因为一次 256 次传递将需要 4 次 64 次传递，因此并行化程度较低。