train.py 文件源码

def train(model, batch, num_samples, word_keep_rate, UNK, alpha):

    xp = model.xp
    use_gpu = (xp == cuda.cupy)

    if use_gpu:
        batch = cuda.to_gpu(batch)

    KL, xents = forward(model, batch, num_samples=num_samples, word_keep_rate=word_keep_rate, UNK=UNK, train=True)
    loss = alpha * KL + sum(xents) / num_samples
    loss.backward()
    optimizer.update()
    loss.unchain_backward()
    if alpha == 0: KL.unchain_backward()