def train_gmm(opt, train_loader, model, board):
model.cuda()
model.train()
# criterion
criterionL1 = nn.L1Loss()
gicloss = GicLoss(opt)
# optimizer
optimizer = torch.optim.Adam(
model.parameters(), lr=opt.lr, betas=(0.5, 0.999))
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda step: 1.0 -
max(0, step - opt.keep_step) / float(opt.decay_step + 1))
for step in range(opt.keep_step + opt.decay_step):
iter_start_time = time.time()
inputs = train_loader.next_batch()
im = inputs['image'].cuda()
im_pose = inputs['pose_image'].cuda()
im_h = inputs['head'].cuda()
shape = inputs['shape'].cuda()
agnostic = inputs['agnostic'].cuda()
c = inputs['cloth'].cuda()
cm = inputs['cloth_mask'].cuda()
im_c = inputs['parse_cloth'].cuda()
im_g = inputs['grid_image'].cuda()
pcm = inputs['parse_cloth_mask'].cuda()
grid, theta = model(agnostic, cm)
warped_cloth = F.grid_sample(c, grid, padding_mode='border')
warped_mask = F.grid_sample(cm, grid, padding_mode='zeros')
warped_grid = F.grid_sample(im_g, grid, padding_mode='zeros')
visuals = [[im_h, shape, im_pose],
[c, warped_cloth, im_c],
[warped_grid, (warped_cloth+im)*0.5, im],
[cm * 2 - 1, warped_mask * 2 - 1, pcm * 2 - 1]]
Lwarp = criterionL1(warped_mask, pcm)
Lgic = gicloss(grid)
Lgic = Lgic / (grid.shape[0] * grid.shape[1] * grid.shape[2])#200x200 = 40.000 * 0.001
loss = 0.0244*Lwarp + 0.9756*Lgic
optimizer.zero_grad()
loss.backward()
optimizer.step()
if (step+1) % opt.display_count == 0:
board_add_images(board, 'combine', visuals, step+1)
board.add_scalar('metric', loss.item(), step + 1)
board.add_scalar('0.9756*Lgic', (0.9756 * Lgic).item(), step + 1)
board.add_scalar('0.0244*Lwarp', (0.0244 * Lwarp).item(), step + 1)
t = time.time() - iter_start_time
print('step: %8d, time: %.3f, loss: %4f, (0.9756*Lgic): %.6f, 0.0244*Lwarp: %.4f' % (step+1, t, loss.item(), (0.9756*Lgic).item(), (0.0244*Lwarp).item()), flush=True)
if (step+1) % opt.save_count == 0:
save_checkpoint(model, os.path.join(
opt.checkpoint_dir, opt.name, 'step_%06d.pth' % (step+1)))
def train_gmm(opt, train_loader, model, board):
model.cuda()
model.train()
# criterion
criterionL1 = nn.L1Loss()
gicloss = GicLoss(opt)
# optimizer
optimizer = torch.optim.Adam(model.parameters(),
lr=opt.lr,
betas=(0.5, 0.999))
scheduler = torch.optim.lr_scheduler.LambdaLR(
optimizer,
lr_lambda=lambda step: 1.0 - max(0, step - opt.keep_step) / float(
opt.decay_step + 1))
for step in range(opt.keep_step + opt.decay_step):
iter_start_time = time.time()
inputs = train_loader.next_batch()
im = inputs['image'].cuda()
im_pose = inputs['pose_image'].cuda()
im_h = inputs['head'].cuda()
shape = inputs['shape'].cuda()
agnostic = inputs['agnostic'].cuda()
c = inputs['cloth'].cuda()
cm = inputs['cloth_mask'].cuda()
im_c = inputs['parse_cloth'].cuda()
im_g = inputs['grid_image'].cuda()
gridtps, thetatps, Ipersp, gridaffine, thetaaffine = model(agnostic, c)
warped_cloth = F.grid_sample(Ipersp, gridtps, padding_mode='border') #
warped_maskaffine = F.grid_sample(cm, gridaffine, padding_mode='zeros')
warped_mask = F.grid_sample(warped_maskaffine,
gridtps,
padding_mode='zeros')
warped_gridaffine = F.grid_sample(im_g,
gridaffine,
padding_mode='zeros')
warped_grid = F.grid_sample(warped_gridaffine,
gridtps,
padding_mode='zeros')
visuals = [[im_h, shape, im_pose], [c, warped_cloth, im_c],
[warped_grid, (warped_cloth * 0.7 + im * 0.3), im],
[Ipersp, (Ipersp * 0.7 + im * 0.3), warped_gridaffine]]
Lwarp = criterionL1(warped_cloth, im_c)
Lpersp = criterionL1(Ipersp, im_c)
LgicTPS = gicloss(gridtps)
LgicTPS = LgicTPS / (gridtps.shape[0] * gridtps.shape[1] *
gridtps.shape[2])
# shape of grid: N, Hout, Wout,2: N x 5 x 5 x 2
# grid shape: N x 5 x 5 x 2
loss = 0.5 * Lpersp + 0.5 * Lwarp # + 40 * LgicTPS
optimizer.zero_grad()
loss.backward()
optimizer.step()
if (step + 1) % opt.display_count == 0:
board_add_images(board, 'combine', visuals, step + 1)
board.add_scalar('metric', loss.item(), step + 1)
board.add_scalar('0.5*Lpersp', (0.5 * Lpersp).item(), step + 1)
#board.add_scalar('40*LgicTPS', (40*LgicTPS).item(), step+1)
board.add_scalar('0.5*Lwarp', (0.5 * Lwarp).item(), step + 1)
t = time.time() - iter_start_time
# print('step: %8d, time: %.3f, loss: %4f, (40*LgicTPS): %.8f, 0.7*Lpersp: %.6f, 0.3*Lwarp: %.6f' % (step+1, t, loss.item(), (40*LgicTPS).item(), (0.7*Lpersp).item(), (0.3*Lwarp).item()), flush=True)
print(
'step: %8d, time: %.3f, loss: %4f, 0.5*Lpersp: %.6f, 0.5*Lwarp: %.6f'
% (step + 1, t, loss.item(), (0.5 * Lpersp).item(),
(0.5 * Lwarp).item()),
flush=True)
if (step + 1) % opt.save_count == 0:
save_checkpoint(
model,
os.path.join(opt.checkpoint_dir, opt.name,
'step_%06d.pth' % (step + 1)))
def train_gmm(opt, train_loader, model, board):
model.cuda()
model.train()
# criterion
criterionL1 = nn.L1Loss()
gicloss = GicLoss(batch_size=opt.batch_size,
imageheight=opt.fine_height,
imagewidth=opt.fine_width)
#criterionVGG = VGGLoss()
# optimizer
optimizer = torch.optim.Adam(model.parameters(),
lr=opt.lr,
betas=(0.5, 0.999))
scheduler = torch.optim.lr_scheduler.LambdaLR(
optimizer,
lr_lambda=lambda step: 1.0 - max(0, step - opt.keep_step) / float(
opt.decay_step + 1))
for step in range(opt.keep_step + opt.decay_step):
iter_start_time = time.time()
inputs = train_loader.next_batch()
im = inputs['image'].cuda()
im_pose = inputs['pose_image'].cuda()
im_h = inputs['head'].cuda()
shape = inputs['shape'].cuda()
agnostic = inputs['agnostic'].cuda()
c = inputs['cloth'].cuda()
cm = inputs['cloth_mask'].cuda()
im_c = inputs['parse_cloth'].cuda()
im_g = inputs['grid_image'].cuda()
grid, theta = model(agnostic, c)
warped_cloth = F.grid_sample(c, grid, padding_mode='border')
warped_mask = F.grid_sample(cm, grid, padding_mode='zeros')
warped_grid = F.grid_sample(im_g, grid, padding_mode='zeros')
visuals = [[im_h, shape, im_pose], [c, warped_cloth, im_c],
[warped_grid, (warped_cloth + im) * 0.5, im]]
# shape of grid: N, Hout, Wout,2: N x 256 x 192 x 2
# grid shape: N x 256 x 192 x 2
#print("grid shape: ", grid.shape)
starttime = time.time()
Gx = grid[:, :, :, 0]
Gy = grid[:, :, :, 1]
dtx = np.zeros([Gx.shape[0], Gx.shape[1], Gx.shape[2]])
dty = np.zeros([Gx.shape[0], Gx.shape[1], Gx.shape[2]])
Gx.cpu()
Gy.cpu()
Lgic = 0
print("start")
for n in range(Gx.shape[0]):
for y in range(0, Gx.shape[1] - 2):
for x in range(0, Gx.shape[2] - 2):
dtx[n, y, x] = DT(Gx[n, y, x].item(), Gy[n, y, x].item(),
Gx[n, y, x + 1].item(), Gy[n, y,
x + 1].item())
dty[n, y, x] = DT(Gx[n, y, x].item(), Gy[n, y, x].item(),
Gx[n, y + 1, x].item(), Gy[n, y + 1,
x].item())
for n in range(Gx.shape[0]):
for y in range(1, Gx.shape[1] - 2):
for x in range(1, Gx.shape[2] - 2):
Lgic = Lgic + np.abs(dtx[n, y - 1, x - 1] - dtx[
n, y - 1, x]) + np.abs(dty[n, y - 1, x - 1] -
dty[n, y, x - 1])
#endtime = time.time()
#print("time:", endtime - starttime)
Lgic = Lgic / (Gx.shape[0] * Gx.shape[1] * Gx.shape[2])
Lwarp = criterionL1(warped_cloth, im_c)
loss = Lwarp + Lgic
#Lgic = gicloss(grid)
endtime = time.time()
print("CalLgic:", endtime - starttime)
#Lwarp = criterionL1(warped_cloth, im_c)
#loss = Lwarp + Lgic
#loss = Lwarp + Lgic
#Lwarp.cuda()
#loss.cuda()
optimizer.zero_grad()
print("begin backward: ", time.time() - endtime)
loss.backward()
print("finish backward: ", time.time() - endtime)
optimizer.step()
if (step + 1) % opt.display_count == 0:
board_add_images(board, 'combine', visuals, step + 1)
board.add_scalar('metric', loss.item(), step + 1)
board.add_scalar('Lgic', Lgic.item(), step + 1)
board.add_scalar('Lwarp', Lwarp.item(), step + 1)
t = time.time() - iter_start_time
print('step: %8d, time: %.3f, loss: %4f, Lgic: %.6f, Lwarp: %.4f' %
(step + 1, t, loss.item(), Lgic.item(), Lwarp.item()),
flush=True)
if (step + 1) % opt.save_count == 0:
save_checkpoint(
model,
os.path.join(opt.checkpoint_dir, opt.name,
'step_%06d.pth' % (step + 1)))
print("finis one step: ", time.time() - endtime)
def train_gmm(opt, train_loader, model, board):
model.cuda()
model.train()
# criterion
criterionL1 = nn.L1Loss()
gicloss = GicLoss(opt)
# optimizer
optimizer = torch.optim.Adam(
model.parameters(), lr=opt.lr, betas=(0.5, 0.999))
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda step: 1.0 -
max(0, step - opt.keep_step) / float(opt.decay_step + 1))
for step in range(opt.keep_step + opt.decay_step):
iter_start_time = time.time()
inputs = train_loader.next_batch()
im = inputs['image'].cuda()
im_pose = inputs['pose_image'].cuda()
im_h = inputs['head'].cuda()
shape = inputs['shape'].cuda()
agnostic = inputs['agnostic'].cuda()
c = inputs['cloth'].cuda()
cm = inputs['cloth_mask'].cuda()
im_c = inputs['parse_cloth'].cuda()
im_g = inputs['grid_image'].cuda()
grid, theta = model(agnostic, cm) # can be added c too for new training
warped_cloth = F.grid_sample(c, grid, padding_mode='border')
warped_mask = F.grid_sample(cm, grid, padding_mode='zeros')
warped_grid = F.grid_sample(im_g, grid, padding_mode='zeros')
visuals = [[im_h, shape, im_pose],
[c, warped_cloth, im_c],
[warped_grid, (warped_cloth+im)*0.5, im]]
# Lwarp = criterionL1(warped_cloth, im_c) # loss for warped cloth
Lwarp = criterionL1(warped_mask, cm) # loss for warped mask thank xuxiaochun025 for fixing the git code.
# grid regularization loss
Lgic = gicloss(grid)
# 200x200 = 40.000 * 0.001
Lgic = Lgic / (grid.shape[0] * grid.shape[1] * grid.shape[2])
loss = Lwarp + 40 * Lgic # total GMM loss
optimizer.zero_grad()
loss.backward()
optimizer.step()
if (step+1) % opt.display_count == 0:
board_add_images(board, 'combine', visuals, step+1)
board.add_scalar('loss', loss.item(), step+1)
board.add_scalar('40*Lgic', (40*Lgic).item(), step+1)
board.add_scalar('Lwarp', Lwarp.item(), step+1)
t = time.time() - iter_start_time
print('step: %8d, time: %.3f, loss: %4f, (40*Lgic): %.8f, Lwarp: %.6f' %
(step+1, t, loss.item(), (40*Lgic).item(), Lwarp.item()), flush=True)
if (step+1) % opt.save_count == 0:
save_checkpoint(model, os.path.join(
opt.checkpoint_dir, opt.name, 'step_%06d.pth' % (step+1)))
def train_gmm(opt, train_loader, model, board):
model.cuda()
model.train()
# criterion
criterionL1 = nn.L1Loss()
gicloss = GicLoss(opt)
#criterionVGG = VGGLoss()
# optimizer
optimizer = torch.optim.Adam(model.parameters(),
lr=opt.lr,
betas=(0.5, 0.999))
scheduler = torch.optim.lr_scheduler.LambdaLR(
optimizer,
lr_lambda=lambda step: 1.0 - max(0, step - opt.keep_step) / float(
opt.decay_step + 1))
for step in range(opt.keep_step + opt.decay_step):
iter_start_time = time.time()
inputs = train_loader.next_batch()
im = inputs['image'].cuda()
im_pose = inputs['pose_image'].cuda()
im_h = inputs['head'].cuda()
shape = inputs['shape'].cuda()
agnostic = inputs['agnostic'].cuda()
c = inputs['cloth'].cuda()
cm = inputs['cloth_mask'].cuda()
im_c = inputs['parse_cloth'].cuda()
im_g = inputs['grid_image'].cuda()
grid, theta = model(agnostic, c)
warped_cloth = F.grid_sample(c, grid, padding_mode='border')
warped_mask = F.grid_sample(cm, grid, padding_mode='zeros')
warped_grid = F.grid_sample(im_g, grid, padding_mode='zeros')
visuals = [[im_h, shape, im_pose], [c, warped_cloth, im_c],
[warped_grid, (warped_cloth + im) * 0.5, im]]
# shape of grid: N, Hout, Wout,2: N x 256 x 192 x 2
# grid shape: N x 256 x 192 x 2
#print("grid shape: ", grid.shape)
#starttime = time.time()
Lgic = gicloss(grid)
#endtime = time.time()
#print("time:", endtime - starttime)
Lgic = Lgic / (grid.shape[0] * grid.shape[1] * grid.shape[2])
Lwarp = criterionL1(warped_cloth, im_c)
loss = Lwarp + Lgic
optimizer.zero_grad()
#print("begin backward: ", time.time() - endtime)
loss.backward()
#print("finish backward: ", time.time() - endtime)
optimizer.step()
if (step + 1) % opt.display_count == 0:
board_add_images(board, 'combine', visuals, step + 1)
board.add_scalar('metric', loss.item(), step + 1)
board.add_scalar('Lgic', Lgic.item(), step + 1)
board.add_scalar('Lwarp', Lwarp.item(), step + 1)
t = time.time() - iter_start_time
print('step: %8d, time: %.3f, loss: %4f, Lgic: %.6f, Lwarp: %.4f' %
(step + 1, t, loss.item(), Lgic.item(), Lwarp.item()),
flush=True)
if (step + 1) % opt.save_count == 0:
save_checkpoint(
model,
os.path.join(opt.checkpoint_dir, opt.name,
'step_%06d.pth' % (step + 1)))
