def train(train_loader, model, optimizer, epoch):
# switch to train mode
model.train()
pbar = tqdm(enumerate(train_loader))
for batch_idx, data in pbar:
data_a, data_p = data
if args.cuda:
data_a, data_p = data_a.float().cuda(), data_p.float().cuda()
data_a, data_p = Variable(data_a), Variable(data_p)
rot_LAFs, inv_rotmat = get_random_rotation_LAFs(data_a, math.pi)
scale = Variable( 0.9 + 0.3* torch.rand(data_a.size(0), 1, 1));
if args.cuda:
scale = scale.cuda()
rot_LAFs[:,0:2,0:2] = rot_LAFs[:,0:2,0:2] * scale.expand(data_a.size(0),2,2)
shift_w, shift_h = get_random_shifts_LAFs(data_a, 2, 2)
rot_LAFs[:,0,2] = rot_LAFs[:,0,2] + shift_w / float(data_a.size(3))
rot_LAFs[:,1,2] = rot_LAFs[:,1,2] + shift_h / float(data_a.size(2))
data_a_rot = extract_patches(data_a, rot_LAFs, PS = data_a.size(2))
st = int((data_p.size(2) - model.PS)/2)
fin = st + model.PS
data_p_crop = data_p[:,:, st:fin, st:fin].contiguous()
data_a_rot_crop = data_a_rot[:,:, st:fin, st:fin].contiguous()
out_a_rot, out_p, out_a = model(data_a_rot_crop,True), model(data_p_crop,True), model(data_a[:,:, st:fin, st:fin].contiguous(), True)
out_p_rotatad = torch.bmm(inv_rotmat, out_p)
######Apply rot and get sifts
out_patches_a_crop = extract_and_crop_patches_by_predicted_transform(data_a_rot, out_a_rot, crop_size = model.PS)
out_patches_p_crop = extract_and_crop_patches_by_predicted_transform(data_p, out_p, crop_size = model.PS)
desc_a = descriptor(out_patches_a_crop)
desc_p = descriptor(out_patches_p_crop)
descr_dist = torch.sqrt(((desc_a - desc_p)**2).view(data_a.size(0),-1).sum(dim=1) + 1e-6).mean()
geom_dist = torch.sqrt(((out_a_rot - out_p_rotatad)**2 ).view(-1,4).sum(dim=1)[0] + 1e-8).mean()
if args.loss == 'HardNet':
loss = loss_HardNet(desc_a,desc_p);
elif args.loss == 'HardNetDetach':
loss = loss_HardNetDetach(desc_a,desc_p);
elif args.loss == 'Geom':
loss = geom_dist;
elif args.loss == 'PosDist':
loss = descr_dist;
else:
print('Unknown loss function')
sys.exit(1)
optimizer.zero_grad()
loss.backward()
optimizer.step()
adjust_learning_rate(optimizer)
if batch_idx % args.log_interval == 0:
pbar.set_description(
'Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.4f}, {:.4f},{:.4f}'.format(
epoch, batch_idx * len(data_a), len(train_loader.dataset),
100. * batch_idx / len(train_loader),
float(loss.detach().cpu().numpy()), float(geom_dist.detach().cpu().numpy()), float(descr_dist.detach().cpu().numpy())))
torch.save({'epoch': epoch + 1, 'state_dict': model.state_dict()},
'{}/checkpoint_{}.pth'.format(LOG_DIR,epoch))
def extract_random_LAF(data, max_rot = math.pi, max_tilt = 1.0, crop_size = 32):
st = int((data.size(2) - crop_size)/2)
fin = st + crop_size
if type(max_rot) is float:
rot_LAFs, inv_rotmat = get_random_rotation_LAFs(data, max_rot)
else:
rot_LAFs = max_rot
inv_rotmat = None
aff_LAFs, inv_TA = get_random_norm_affine_LAFs(data, max_tilt);
aff_LAFs[:,0:2,0:2] = torch.bmm(rot_LAFs[:,0:2,0:2],aff_LAFs[:,0:2,0:2])
data_aff = extract_patches(data, aff_LAFs, PS = data.size(2))
data_affcrop = data_aff[:,:, st:fin, st:fin].contiguous()
return data_affcrop, data_aff, rot_LAFs,inv_rotmat,inv_TA
def test(test_loader, model, epoch):
# switch to evaluate mode
model.eval()
geom_distances, desc_distances = [], []
pbar = tqdm(enumerate(test_loader))
for batch_idx, (data_a, data_p) in pbar:
if args.cuda:
data_a, data_p = data_a.float().cuda(), data_p.float().cuda()
data_a, data_p = Variable(data_a, volatile=True), Variable(data_p, volatile=True)
rot_LAFs, inv_rotmat = get_random_rotation_LAFs(data_a, math.pi)
data_a_rot = extract_patches(data_a, rot_LAFs, PS = data_a.size(2))
st = int((data_p.size(2) - model.PS)/2)
fin = st + model.PS
data_p = data_p[:,:, st:fin, st:fin].contiguous()
data_a_rot = data_a_rot[:,:, st:fin, st:fin].contiguous()
out_a_rot, out_p = model(data_a_rot, True), model(data_p, True)
out_p_rotatad = torch.bmm(inv_rotmat, out_p)
geom_dist = torch.sqrt((out_a_rot - out_p_rotatad)**2 + 1e-12).mean()
out_patches_a_crop = extract_and_crop_patches_by_predicted_transform(data_a_rot, out_a_rot, crop_size = model.PS)
out_patches_p_crop = extract_and_crop_patches_by_predicted_transform(data_p, out_p, crop_size = model.PS)
desc_a = descriptor(out_patches_a_crop)
desc_p = descriptor(out_patches_p_crop)
descr_dist = torch.sqrt(((desc_a - desc_p)**2).view(data_a.size(0),-1).sum(dim=1) + 1e-6)#/ float(desc_a.size(1))
descr_dist = descr_dist.mean()
geom_distances.append(geom_dist.data.cpu().numpy().reshape(-1,1))
desc_distances.append(descr_dist.data.cpu().numpy().reshape(-1,1))
if batch_idx % args.log_interval == 0:
pbar.set_description(' Test Epoch: {} [{}/{} ({:.0f}%)]'.format(
epoch, batch_idx * len(data_a), len(test_loader.dataset),
100. * batch_idx / len(test_loader)))
geom_distances = np.vstack(geom_distances).reshape(-1,1)
desc_distances = np.vstack(desc_distances).reshape(-1,1)
print('\33[91mTest set: Geom MSE: {:.8f}\n\33[0m'.format(geom_distances.mean()))
print('\33[91mTest set: Desc dist: {:.8f}\n\33[0m'.format(desc_distances.mean()))
return
def train(train_loader, model, optimizer, epoch):
# switch to train mode
model.train()
pbar = tqdm(enumerate(train_loader))
for batch_idx, data in pbar:
data_a, data_p = data
if args.cuda:
data_a, data_p = data_a.float().cuda(), data_p.float().cuda()
data_a, data_p = Variable(data_a), Variable(data_p)
st = int((data_p.size(2) - model.PS) / 2)
fin = st + model.PS
#
#
max_tilt = 3.0
if epoch > 1:
max_tilt = 4.0
if epoch > 3:
max_tilt = 4.5
if epoch > 5:
max_tilt = 4.8
rot_LAFs_a, inv_rotmat_a = get_random_rotation_LAFs(data_a, math.pi)
aff_LAFs_a, inv_TA_a = get_random_norm_affine_LAFs(data_a, max_tilt)
aff_LAFs_a[:, 0:2, 0:2] = torch.bmm(rot_LAFs_a[:, 0:2, 0:2],
aff_LAFs_a[:, 0:2, 0:2])
data_a_aff = extract_patches(data_a, aff_LAFs_a, PS=data_a.size(2))
data_a_aff_crop = data_a_aff[:, :, st:fin, st:fin].contiguous()
aff_LAFs_p, inv_TA_p = get_random_norm_affine_LAFs(data_p, max_tilt)
aff_LAFs_p[:, 0:2, 0:2] = torch.bmm(rot_LAFs_a[:, 0:2, 0:2],
aff_LAFs_p[:, 0:2, 0:2])
data_p_aff = extract_patches(data_p, aff_LAFs_p, PS=data_p.size(2))
data_p_aff_crop = data_p_aff[:, :, st:fin, st:fin].contiguous()
out_a_aff, out_p_aff = model(data_a_aff_crop,
True), model(data_p_aff_crop, True)
out_p_aff_back = torch.bmm(inv_TA_p, out_p_aff)
out_a_aff_back = torch.bmm(inv_TA_a, out_a_aff)
######Apply rot and get sifts
out_patches_a_crop = extract_and_crop_patches_by_predicted_transform(
data_a_aff, out_a_aff, crop_size=model.PS)
out_patches_p_crop = extract_and_crop_patches_by_predicted_transform(
data_p_aff, out_p_aff, crop_size=model.PS)
desc_a = descriptor(out_patches_a_crop)
desc_p = descriptor(out_patches_p_crop)
descr_dist = torch.sqrt((
(desc_a - desc_p)**2).view(data_a.size(0), -1).sum(dim=1) + 1e-6)
descr_loss = loss_HardNet(desc_a, desc_p, anchor_swap=True)
geom_dist = torch.sqrt((
(out_a_aff_back - out_p_aff_back)**2).view(-1, 4).mean(dim=1) +
1e-8)
if args.merge == 'sum':
loss = descr_loss
elif args.merge == 'mul':
loss = descr_loss
else:
print('Unknown merge option')
sys.exit(0)
optimizer.zero_grad()
loss.backward()
optimizer.step()
adjust_learning_rate(optimizer)
if batch_idx % 2 == 0:
pbar.set_description(
'Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.4f}, {},{:.4f}'.
format(epoch, batch_idx * len(data_a),
len(train_loader.dataset),
100. * batch_idx / len(train_loader), loss.data[0],
geom_dist.mean().data[0],
descr_dist.mean().data[0]))
torch.save({
'epoch': epoch + 1,
'state_dict': model.state_dict()
}, '{}/checkpoint_{}.pth'.format(LOG_DIR, epoch))