CUDA批出memory

Question

我有一个小数据集并运行一个名为 LightXML 的脚本，该脚本位于 git 上： https://github.com/kongds/LightXML

我收到此错误：

torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 192.00 MiB (GPU 0; 6.00 GiB total capacity; 4.71 GiB already allocated; 0 bytes free; 4.82 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

我发现有多个帖子说我应该减少批量大小，但我似乎找不到它的定义。

脚本停止在：

train_loss = model.one_epoch(epoch, trainloader, optimizer, mode='train',
                                 eval_loader=validloader if args.valid else testloader,
                                 eval_step=args.eval_step, log=LOG)

one_epoch 方法如下，脚本停滞在outputs = self(**inputs)

def one_epoch(self, epoch, dataloader, optimizer,
              mode='train', eval_loader=None, eval_step=20000, log=None):

    bar = tqdm.tqdm(total=len(dataloader))
    p1, p3, p5 = 0, 0, 0
    g_p1, g_p3, g_p5 = 0, 0, 0
    total, acc1, acc3, acc5 = 0, 0, 0, 0
    g_acc1, g_acc3, g_acc5 = 0, 0, 0
    train_loss = 0

    if mode == 'train':
        self.train()
    else:
        self.eval()

    if self.use_swa and epoch == self.swa_warmup_epoch and mode == 'train':
        self.swa_init()

    if self.use_swa and mode == 'eval':
        self.swa_swap_params()

    pred_scores, pred_labels = [], []
    bar.set_description(f'{mode}-{epoch}')

    with torch.set_grad_enabled(mode == 'train'):
        for step, data in enumerate(dataloader):
            batch = tuple(t for t in data)
            have_group = len(batch) > 4

            check_memory()
            inputs = {'input_ids':      batch[0].cuda(),
                      'attention_mask': batch[1].cuda(),
                      'token_type_ids': batch[2].cuda()}


            if mode == 'train':
                inputs['labels'] = batch[3].cuda()

                if self.group_y is not None:
                    inputs['group_labels'] = batch[4].cuda()
                    inputs['candidates'] = batch[5].cuda()
            print("------------------------------------------------------------")
            outputs = self(**inputs)
            bar.update(1)


            if mode == 'train':
                loss = outputs[1]
                loss /= self.update_count
                train_loss += loss.item()

                with amp.scale_loss(loss, optimizer) as scaled_loss:
                    scaled_loss.backward()

                if step % self.update_count == 0:
                    optimizer.step()
                    self.zero_grad()

                if step % eval_step == 0 and eval_loader is not None and step != 0:
                    results = self.one_epoch(epoch, eval_loader, optimizer, mode='eval')
                    p1, p3, p5 = results[3:6]
                    g_p1, g_p3, g_p5 = results[:3]
                    if self.group_y is not None:
                        log.log(f'{epoch:>2} {step:>6}: {p1:.4f}, {p3:.4f}, {p5:.4f}'
                                f' {g_p1:.4f}, {g_p3:.4f}, {g_p5:.4f}')
                    else:
                        log.log(f'{epoch:>2} {step:>6}: {p1:.4f}, {p3:.4f}, {p5:.4f}')
                    # NOTE: we don't reset model to train mode and keep model in eval mode
                    # which means all dropout will be remove after `eval_step` in every epoch
                    # this tricks makes LightXML converge fast
                    # self.train()

                if self.use_swa and step % self.swa_update_step == 0:
                    self.swa_step()

                bar.set_postfix(loss=loss.item())
            elif self.group_y is None:
                logits = outputs
                if mode == 'eval':
                    labels = batch[3]
                    _total, _acc1, _acc3, _acc5 =  self.get_accuracy(None, logits, labels.cpu().numpy())
                    total += _total; acc1 += _acc1; acc3 += _acc3; acc5 += _acc5
                    p1 = acc1 / total
                    p3 = acc3 / total / 3
                    p5 = acc5 / total / 5
                    bar.set_postfix(p1=p1, p3=p3, p5=p5)
                elif mode == 'test':
                    pred_scores.append(logits.detach().cpu())
            else:
                group_logits, candidates, logits = outputs

                if mode == 'eval':
                    labels = batch[3]
                    group_labels = batch[4]

                    _total, _acc1, _acc3, _acc5 = self.get_accuracy(candidates, logits, labels.cpu().numpy())
                    total += _total; acc1 += _acc1; acc3 += _acc3; acc5 += _acc5
                    p1 = acc1 / total
                    p3 = acc3 / total / 3
                    p5 = acc5 / total / 5

                    _, _g_acc1, _g_acc3, _g_acc5 = self.get_accuracy(None, group_logits, group_labels.cpu().numpy())
                    g_acc1 += _g_acc1; g_acc3 += _g_acc3; g_acc5 += _g_acc5
                    g_p1 = g_acc1 / total
                    g_p3 = g_acc3 / total / 3
                    g_p5 = g_acc5 / total / 5
                    bar.set_postfix(p1=p1, p3=p3, p5=p5, g_p1=g_p1, g_p3=g_p3, g_p5=g_p5)
                elif mode == 'test':
                    _scores, _indices = torch.topk(logits.detach().cpu(), k=100)
                    _labels = torch.stack([candidates[i][_indices[i]] for i in range(_indices.shape[0])], dim=0)
                    pred_scores.append(_scores.cpu())
                    pred_labels.append(_labels.cpu())


    if self.use_swa and mode == 'eval':
        self.swa_swap_params()
    bar.close()

    if mode == 'eval':
        return g_p1, g_p3, g_p5, p1, p3, p5
    elif mode == 'test':
        return torch.cat(pred_scores, dim=0).numpy(), torch.cat(pred_labels, dim=0).numpy() if len(pred_labels) != 0 else None
    elif mode == 'train':
        return train_loss

我还包括在 outs=self.bert() 处停止的转发 function

def forward(self, input_ids, attention_mask, token_type_ids,
            labels=None, group_labels=None, candidates=None):
    is_training = labels is not None
   
    outs = self.bert(
        input_ids,
        attention_mask=attention_mask,
        token_type_ids=token_type_ids
    )[-1]

    out = torch.cat([outs[-i][:, 0] for i in range(1, self.feature_layers+1)], dim=-1)

    out = self.drop_out(out)
    group_logits = self.l0(out)

    if self.group_y is None:
        logits = group_logits
        if is_training:
            loss_fn = torch.nn.BCEWithLogitsLoss()
            loss = loss_fn(logits, labels)
            return logits, loss
        else:
            return logits

    if is_training:
        l = labels.to(dtype=torch.bool)
        target_candidates = torch.masked_select(candidates, l).detach().cpu()
        target_candidates_num = l.sum(dim=1).detach().cpu()
    groups, candidates, group_candidates_scores = self.get_candidates(group_logits,
                                                                      group_gd=group_labels if is_training else None)
    if is_training:
        bs = 0
        new_labels = []
        for i, n in enumerate(target_candidates_num.numpy()):
            be = bs + n
            c = set(target_candidates[bs: be].numpy())
            c2 = candidates[i]
            new_labels.append(torch.tensor([1.0 if i in c else 0.0 for i in c2 ]))
            if len(c) != new_labels[-1].sum():
                s_c2 = set(c2)
                for cc in list(c):
                    if cc in s_c2:
                        continue
                    for j in range(new_labels[-1].shape[0]):
                        if new_labels[-1][j].item() != 1:
                            c2[j] = cc
                            new_labels[-1][j] = 1.0
                            break
            bs = be
        labels = torch.stack(new_labels).cuda()
    candidates, group_candidates_scores =  torch.LongTensor(candidates).cuda(), torch.Tensor(group_candidates_scores).cuda()

    emb = self.l1(out)
    embed_weights = self.embed(candidates) # N, sampled_size, H
    emb = emb.unsqueeze(-1)
    logits = torch.bmm(embed_weights, emb).squeeze(-1)

    if is_training:
        loss_fn = torch.nn.BCEWithLogitsLoss()
        loss = loss_fn(logits, labels) + loss_fn(group_logits, group_labels)
        return logits, loss
    else:
        candidates_scores = torch.sigmoid(logits)
        candidates_scores = candidates_scores * group_candidates_scores
        return group_logits, candidates, candidates_scores

Answer 1

是的，问题可能出在 batch_size 上。 通常 batch_size 定义在 DataLoader 中。 在main.py中可以看到：

trainloader = DataLoader(train_d, batch_size=args.batch, num_workers=5,
                                 shuffle=True)

您可以在脚本运行时定义它 - args.batch。