natanielruiz
2017-08-12 3de9cc574450403fc33e9dfd4ae298d219e2ea95
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.'