Получите LR от скорости циклического обучения в PyTorch

Question

Я пытаюсь реализовать подход циклическая скорость обучения поверх PyTorch повторной реализации StyleGAN с помощью rosinality . Для этого я опираюсь на то, что было предложено в этом сообщении в блоге.

Чтобы проверить, как изменяется потеря в зависимости от скорости обучения, мне нужно построить график того, как (LR, потеря) изменения. Здесь вы можете найти мою модифицированную версию train.py. Вот основные изменения, которые я сделал:

1. Define cyclical_lr, функция, регулирующая циклическую скорость обучения

def cyclical_lr(stepsize, min_lr, max_lr):

    # Scaler: we can adapt this if we do not want the triangular CLR
    scaler = lambda x: 1.

    # Lambda function to calculate the LR
    lr_lambda = lambda it: min_lr + (max_lr - min_lr) * relative(it, stepsize)

    # Additional function to see where on the cycle we are
    def relative(it, stepsize):
        cycle = math.floor(1 + it / (2 * stepsize))
        x = abs(it / stepsize - 2 * cycle + 1)
        return max(0, (1 - x)) * scaler(cycle)

    return lr_lambda

Реализация скорости циклического обучения как для дискриминатора, так и для генератора.

step_size = 5*256
end_lr = 10**-1
factor = 10**5
clr = cyclical_lr(step_size, min_lr=end_lr / factor, max_lr=end_lr)
scheduler_g = torch.optim.lr_scheduler.LambdaLR(g_optimizer, [clr, clr])

d_optimizer = optim.Adam(discriminator.parameters(), lr=args.lr, betas=(0.0, 0.99))
scheduler_d = torch.optim.lr_scheduler.LambdaLR(d_optimizer, [clr])

У вас есть предложения о том, как построить график изменения потерь в зависимости от скорость обучения? В идеале я хотел бы сделать это с помощью TensorBoard, где на данный момент я рисую потери генератора, потери дискриминатора и размер сгенерированных изображений в зависимости от числа итераций:

    if (i + 1) % 100 == 0:
        writer.add_scalar('Loss/G', gen_loss_val, i * args.batch.get(resolution))
        writer.add_scalar('Loss/D', disc_loss_val, i * args.batch.get(resolution))
        writer.add_scalar('Step/pixel_size', (4 * 2 ** step), i * args.batch.get(resolution))
        print(args.batch.get(resolution))