AFLW training ready.

natanielruiz

2017-08-12 beb9f36419d0df03c3248757f54af032a633e05c

 code/train_shape.py

@@ -128,13 +128,12 @@
    reg_criterion = nn.MSELoss().cuda(gpu)
    # Regression loss coefficient
    alpha = 0.1
    lsm = nn.Softmax()

    idx_tensor = [idx for idx in xrange(66)]
    idx_tensor = torch.FloatTensor(idx_tensor).cuda(gpu)

    optimizer = torch.optim.Adam([{'params': get_ignored_params(model), 'lr': args.lr},
                                  {'params': get_non_ignored_params(model), 'lr': args.lr}],
                                  {'params': get_non_ignored_params(model), 'lr': args.lr * 10}],
                                  lr = args.lr)

    print 'Ready to train network.'