首页 > 解决方案 > 无需教师强制的高效变压器培训

问题描述

我想在不依赖教师强制的情况下训练一个用于时间序列预测的变压器模型。这是我的训练循环:

# src shape is [n_batch, src_length, d_model] --> [N, S, D]
# tgt and gt (groundtruth) shape is [n_batch, tgt_length, d_model] --> [N, T, D]

l = 0.
for src, tgt, gt in loader:
    src, tgt, gt = src.to(device), tgt.to(device), gt.to(device)
    optimizer.zero_grad()

    # first tgt element is the last src one, followed by a zero-filled tensor
    tgt = torch.cat((src[:, -1:, :], torch.zeros(tgt.shape[0], tgt.shape[1] - 1, tgt.shape[2], device=device)), dim=1)

    out = torch.zeros_like(tgt, device=device)
    memory = model.encode(src)
    for j in range(gt.shape[1]):
        out = model.decode(tgt, memory)
        l += loss(out[:, j, :], gt[:, j, :])
        tgt = tgt.detach()
        tgt = torch.cat((tgt[:, :j + 1, :], out[:, j:j+1, :], tgt[:, j + 1:gt.shape[1] - 1, :]), dim=1) 
    l.backward()

我开始用源序列的最后一个元素初始化目标,后面跟着尽可能多的零来匹配正确的形状。然后,在内部for 中,我将新的预测值附加到前一个目标上(再次添加所需的零)以创建下一次迭代的目标。

首先,我想确保我的代码是否正确,至少从概念的角度来看是这样。

现在,尽管单个 epoch 的训练时间增加了 10 倍,但上面的代码对于非常小的批次(数十个样本)运行良好。使用教师强迫训练相同的模型。问题是,如果我尝试使用“正常”批量大小,比如 32、64 个样本,代码会在训练期间崩溃,因为 CUDA 内存不足。我的数据集的每个样本的形状都是 [12, 5],在嵌入到模型中之后变成了 [12, 128]。

我的本地机器 GPU 只有 4GB 的 RAM,但使用 Colab 的 12Gbs GPU 时也会发生相同的行为。我很确定有一种更好的方法可以在没有教师强迫的情况下实施变压器培训。

RuntimeError: CUDA out of memory. Tried to allocate 2.00 MiB (GPU 0; 11.17 GiB total capacity; 10.38 GiB already allocated; 1.81 MiB free; 10.67 GiB reserved in total by PyTorch)

上面的错误通常在调用model.encode(src)model.decode(tgt, memory)时验证。另一方面,使用教师强制代码工作正常。

我之前没有利用torch.cat函数,而是尝试简单地创建一个新的空张量并将当前tgt加上每个时间步的新预测输出。不幸的是,autograd不喜欢这样,它抛出了关于正在就地编辑的变量的错误

标签: pythonpytorchtime-seriestransformer

解决方案


推荐阅读