def validate(net, curr_iter, optimizer):
print('validating...')
net.eval()
loss_record1 = AvgMeter()
iter_num1 = len(test1_loader)
with torch.no_grad():
for i, data in enumerate(test1_loader):
inputs, gts = data
inputs = Variable(inputs).cuda()
gts = Variable(gts).cuda()
res = net(inputs)
loss = criterion(res, gts)
loss_record1.update(loss.data, inputs.size(0))
print('processed test1 %d / %d' % (i + 1, iter_num1))
# snapshot_name = 'iter_%d_loss1_%.5f_lr_%.6f' % (curr_iter + 1, loss_record1.avg,
# optimizer.param_groups[1]['lr'])
log_val = '[validate]: [iter %d], [loss1 %.5f]' % (curr_iter + 1,
loss_record1.avg)
print(log_val)
open(log_path_val, 'a').write(log_val + '\n')
# torch.save(net.state_dict(), os.path.join(ckpt_path, exp_name, snapshot_name + '.pth'))
# torch.save(optimizer.state_dict(), os.path.join(ckpt_path, exp_name, snapshot_name + '_optim.pth'))
net.train()
def main():
net = R3Net().cuda()
print 'load snapshot \'%s\' for testing' % args['snapshot']
net.load_state_dict(
torch.load(os.path.join(ckpt_path, exp_name,
args['snapshot'] + '.pth')))
net.eval()
results = {}
with torch.no_grad():
for name, root in to_test.iteritems():
precision_record, recall_record, = [
AvgMeter() for _ in range(256)
], [AvgMeter() for _ in range(256)]
mae_record = AvgMeter()
if args['save_results']:
check_mkdir(
os.path.join(
ckpt_path, exp_name,
'(%s) %s_%s' % (exp_name, name, args['snapshot'])))
img_list = os.listdir(root)
print img_list
for idx, img_name in enumerate(img_list):
print 'predicting for %s: %d / %d' % (name, idx + 1,
len(img_list))
#print img_name
img_path = os.path.join(root, img_name)
img = Image.open(img_path).convert('RGB')
img_var = Variable(img_transform(img).unsqueeze(0),
volatile=True).cuda()
prediction = net(img_var)
prediction = np.array(to_pil(prediction.data.squeeze(0).cpu()))
if args['crf_refine']:
prediction = crf_refine(np.array(img), prediction)
#gt = np.array(Image.open(os.path.join(root+"/masks", img_name )).convert('L'))
#precision, recall, mae = cal_precision_recall_mae(prediction, gt)
#for pidx, pdata in enumerate(zip(precision, recall)):
# p, r = pdata
# precision_record[pidx].update(p)
# recall_record[pidx].update(r)
#mae_record.update(mae)
if args['save_results']:
Image.fromarray(prediction).save(
os.path.join(
ckpt_path, exp_name, '%s_%s_%s' %
(exp_name, name, args['snapshot'] + "kaist_test"),
img_name))
def train(net, optimizer):
curr_iter = args['last_iter']
while True:
total_loss_record, loss0_record, loss1_record, loss2_record = AvgMeter(), AvgMeter(), AvgMeter(), AvgMeter()
loss3_record, loss4_record, loss5_record, loss6_record = AvgMeter(), AvgMeter(), AvgMeter(), AvgMeter()
for i, data in enumerate(train_loader):
optimizer.param_groups[0]['lr'] = 2 * args['lr'] * (1 - float(curr_iter) / args['iter_num']
) ** args['lr_decay']
optimizer.param_groups[1]['lr'] = args['lr'] * (1 - float(curr_iter) / args['iter_num']
) ** args['lr_decay']
inputs, labels = data
batch_size = inputs.size(0)
inputs = Variable(inputs).cuda()
labels = Variable(labels).cuda()
optimizer.zero_grad()
outputs0, outputs1, outputs2, outputs3, outputs4, outputs5, outputs6 = net(inputs)
loss0 = criterion(outputs0, labels)
loss1 = criterion(outputs1, labels)
loss2 = criterion(outputs2, labels)
loss3 = criterion(outputs3, labels)
loss4 = criterion(outputs4, labels)
loss5 = criterion(outputs5, labels)
loss6 = criterion(outputs6, labels)
total_loss = loss0 + loss1 + loss2 + loss3 + loss4 + loss5 + loss6
total_loss.backward()
optimizer.step()
total_loss_record.update(total_loss.data[0], batch_size)
loss0_record.update(loss0.data[0], batch_size)
loss1_record.update(loss1.data[0], batch_size)
loss2_record.update(loss2.data[0], batch_size)
loss3_record.update(loss3.data[0], batch_size)
loss4_record.update(loss4.data[0], batch_size)
loss5_record.update(loss5.data[0], batch_size)
loss6_record.update(loss6.data[0], batch_size)
curr_iter += 1
log = '[iter %d], [total loss %.5f], [loss0 %.5f], [loss1 %.5f], [loss2 %.5f], [loss3 %.5f], ' \
'[loss4 %.5f], [loss5 %.5f], [loss6 %.5f], [lr %.13f]' % \
(curr_iter, total_loss_record.avg, loss0_record.avg, loss1_record.avg, loss2_record.avg,
loss3_record.avg, loss4_record.avg, loss5_record.avg, loss6_record.avg,
optimizer.param_groups[1]['lr'])
print
log
open(log_path, 'a').write(log + '\n')
if curr_iter == args['iter_num']:
torch.save(net.state_dict(), os.path.join(ckpt_path, exp_name, '%d.pth' % curr_iter))
torch.save(optimizer.state_dict(),
os.path.join(ckpt_path, exp_name, '%d_optim.pth' % curr_iter))
return
def train(net, optimizer):
curr_iter = args['last_iter']
while True:
total_loss_record, loss0_record, loss1_record, loss2_record = AvgMeter(), AvgMeter(), AvgMeter(), AvgMeter()
for i, data in enumerate(train_loader):
optimizer.param_groups[0]['lr'] = 2 * args['lr'] * (1 - float(curr_iter) / args['iter_num']
) ** args['lr_decay']
optimizer.param_groups[1]['lr'] = args['lr'] * (1 - float(curr_iter) / args['iter_num']
) ** args['lr_decay']
inputs, depth,labels,img_name = data
batch_size = inputs.size(0)
inputs = Variable(inputs).cuda()
depth = Variable(depth).cuda()
labels = Variable(labels).cuda()
optimizer.zero_grad()
depthMap, salMap, sideout5, sideout4, sideout3, sideout2 = net(inputs, depth)
loss0 = criterion(depthMap, labels)
loss1 = criterion(salMap, labels)
loss5 = criterion(sideout5, labels)
loss4 = criterion(sideout4, labels)
loss3 = criterion(sideout3, labels)
loss2 = criterion(sideout2, labels)
total_loss = loss0 + loss1 + loss5 + loss4 + loss3 + loss2
total_loss.backward()
optimizer.step()
total_loss_record.update(total_loss.item(), batch_size)
loss0_record.update(loss0.item(), batch_size)
loss1_record.update(loss1.item(), batch_size)
curr_iter += 1
log = '[iter %d], [total loss %.5f], [loss0 %.5f],[loss1 %.5f],[lr %.13f]' % \
(curr_iter, total_loss_record.avg, loss0_record.avg, loss1_record.avg,
optimizer.param_groups[1]['lr'])
print(log)
open(log_path, 'a').write(log + '\n')
if curr_iter == args['iter_num']:
torch.save(net.state_dict(), os.path.join(ckpt_path, exp_name, '%d_res18_noCross_Gdept.pth' % curr_iter))
return
def train(net, optimizer):
curr_iter = args['last_iter']
while True:
total_loss_record, loss1_record, loss2_record, loss3_record, loss4_record, loss5_record, loss6_record, loss7_record, loss8_record = AvgMeter(), AvgMeter(), AvgMeter(), AvgMeter(), AvgMeter(), AvgMeter(), AvgMeter(), AvgMeter(), AvgMeter()
for i, data in enumerate(train_loader):
optimizer.param_groups[0]['lr'] = 2 * args['lr'] * (1 - float(curr_iter) / args['iter_num']
) ** args['lr_decay']
optimizer.param_groups[1]['lr'] = args['lr'] * (1 - float(curr_iter) / args['iter_num']
) ** args['lr_decay']
# data\binarizing\Variable
inputs, labels = data
labels[labels > 0.5] = 1
labels[labels != 1] = 0
batch_size = inputs.size(0)
inputs = Variable(inputs).cuda()
labels = Variable(labels).cuda()
optimizer.zero_grad()
output_fpn, output_final = net(inputs)
##########loss#############
loss1 = criterion(output_fpn, labels)
loss2 = criterion(output_final, labels)
total_loss = loss1+loss2
total_loss.backward()
optimizer.step()
total_loss_record.update(total_loss.item(), batch_size)
loss1_record.update(loss1.item(), batch_size)
loss2_record.update(loss2.item(), batch_size)
#############log###############
curr_iter += 1
log = '[iter %d], [total loss %.5f],[loss1 %.5f],[loss1 %.5f],[lr %.13f] ' % \
(curr_iter, total_loss_record.avg, loss1_record.avg, loss2_record.avg, optimizer.param_groups[1]['lr'])
print(log)
open(log_path, 'a').write(log + '\n')
if curr_iter == args['iter_num']:
torch.save(net.state_dict(), os.path.join(ckpt_path, exp_name, '%d.pth' % curr_iter))
torch.save(optimizer.state_dict(),
os.path.join(ckpt_path, exp_name, '%d_optim.pth' % curr_iter))
return
def train(epoch):
print('\nEpoch: %d' % epoch)
net.train()
train_loss_record, bce_loss_record, dice_loss_record = AvgMeter(
), AvgMeter(), AvgMeter()
for batch_idx, data in enumerate(train_loader):
if epoch == args['lr_step']:
optimizer.param_groups[0]['lr'] = 2 * args['lr'] / args['lr_decay']
optimizer.param_groups[1]['lr'] = args['lr'] / args['lr_decay']
inputs, labels, counter = data
batch_size = inputs.size(0)
inputs = Variable(inputs).cuda()
labels = Variable(labels).cuda()
counter = Variable(counter).cuda()
optimizer.zero_grad()
outputs, outputs_counter = net(inputs)
# outputs = net(inputs)
# BCE loss and dice loss can be used
criterion_bce = nn.BCELoss()
criterion_dice = Dice_loss()
# if not isinstance(fnd_out, list):
loss_bce = criterion_bce(outputs, labels) + criterion_bce(
outputs_counter, counter)
loss_dice = criterion_dice(outputs, labels) + criterion_dice(
outputs_counter, counter)
# loss_bce = criterion_bce(outputs, labels)
# loss_dice = criterion_dice(outputs, labels)
# else:
# loss_bce = criterion_bce(outputs, labels)
# loss_dice = criterion_dice(outputs, labels)
# for fnd_mask, fpd_mask in zip(fnd_out, fpd_out):
# loss_bce += criterion_bce(fnd_mask, fnd) + criterion_bce(fpd_mask, fpd)
# else:
# loss_bce_each = [None] * len(outputs)
# loss_dice_each = [None] * len(outputs)
# for idx in range(len(outputs)):
# loss_bce_each[idx] = criterion_bce(outputs[idx], labels)
# loss_dice_each[idx] = criterion_dice(outputs[idx], labels)
# loss_bce = sum(loss_bce_each)
# loss_dice = sum(loss_dice_each)
coeff = loss_dice.item() / loss_bce.item(
) if loss_dice.item() / loss_bce.item() < 1 else 1
# coeff = 1
loss = coeff * loss_bce + loss_dice
# loss = loss_bce + loss_dice
loss.backward()
optimizer.step()
train_loss_record.update(loss.item(), batch_size)
bce_loss_record.update(loss_bce.item(), batch_size)
dice_loss_record.update(loss_dice.item(), batch_size)
log = 'iter: %d | [bce loss: %.5f], [dice loss: %.5f],[Total loss: %.5f], [lr: %.8f]' % \
(epoch, bce_loss_record.avg, dice_loss_record.avg, train_loss_record.avg, optimizer.param_groups[1]['lr'])
progress_bar(batch_idx, len(train_loader), log)
def main():
net = R3Net().cuda()
print
'load snapshot \'%s\' for testing' % args['snapshot']
net.load_state_dict(torch.load(os.path.join(ckpt_path, exp_name, args['snapshot'] + '.pth')))
net.eval()
results = {}
for name, root in to_test.iteritems():
precision_record, recall_record, = [AvgMeter() for _ in range(256)], [AvgMeter() for _ in range(256)]
mae_record = AvgMeter()
img_list = [os.path.splitext(f)[0] for f in os.listdir(root) if f.endswith('.jpg')]
for idx, img_name in enumerate(img_list):
print
'predicting for %s: %d / %d' % (name, idx + 1, len(img_list))
check_mkdir(os.path.join(ckpt_path, exp_name, '(%s) %s_%s' % (exp_name, name, args['snapshot'])))
img = Image.open(os.path.join(root, img_name + '.jpg')).convert('RGB')
img_var = Variable(img_transform(img).unsqueeze(0), volatile=True).cuda()
prediction = net(img_var)
prediction = np.array(to_pil(prediction.data.squeeze(0).cpu()))
if args['crf_refine']:
prediction = crf_refine(np.array(img), prediction)
gt = np.array(Image.open(os.path.join(root, img_name + '.png')).convert('L'))
precision, recall, mae = cal_precision_recall_mae(prediction, gt)
for pidx, pdata in enumerate(zip(precision, recall)):
p, r = pdata
precision_record[pidx].update(p)
recall_record[pidx].update(r)
mae_record.update(mae)
if args['save_results']:
Image.fromarray(prediction).save(os.path.join(ckpt_path, exp_name, '(%s) %s_%s' % (
exp_name, name, args['snapshot']), img_name + '.png'))
fmeasure = cal_fmeasure([precord.avg for precord in precision_record],
[rrecord.avg for rrecord in recall_record])
results[name] = {'fmeasure': fmeasure, 'mae': mae_record.avg}
print
'test results:'
print
results
def main():
net = Ensemble(device_id, pretrained=False)
print ('load snapshot \'%s\' for testing' % args['snapshot'])
# net.load_state_dict(torch.load('pretrained/R2Net.pth', map_location='cuda:2'))
# net = load_part_of_model2(net, 'pretrained/R2Net.pth', device_id=2)
net.load_state_dict(torch.load(os.path.join(ckpt_path, exp_name, args['snapshot'] + '.pth'),
map_location='cuda:' + str(device_id)))
net.eval()
net.cuda()
results = {}
with torch.no_grad():
for name, root in to_test.items():
precision_record, recall_record, = [AvgMeter() for _ in range(256)], [AvgMeter() for _ in range(256)]
mae_record = AvgMeter()
if args['save_results']:
check_mkdir(os.path.join(ckpt_path, exp_name, '(%s) %s_%s' % (exp_name, name, args['snapshot'])))
img_list = [i_id.strip() for i_id in open(imgs_path)]
for idx, img_name in enumerate(img_list):
print('predicting for %s: %d / %d' % (name, idx + 1, len(img_list)))
print(img_name)
if name == 'VOS' or name == 'DAVSOD':
img = Image.open(os.path.join(root, img_name + '.png')).convert('RGB')
else:
img = Image.open(os.path.join(root, img_name + '.jpg')).convert('RGB')
shape = img.size
img = img.resize(args['input_size'])
img_var = Variable(img_transform(img).unsqueeze(0), volatile=True).cuda()
start = time.time()
outputs_a, outputs_c = net(img_var)
a_out1u, a_out2u, a_out2r, a_out3r, a_out4r, a_out5r = outputs_a # F3Net
# b_outputs0, b_outputs1 = outputs_b # CPD
c_outputs0, c_outputs1, c_outputs2, c_outputs3, c_outputs4 = outputs_c # RAS
prediction = torch.sigmoid(c_outputs0)
end = time.time()
print('running time:', (end - start))
# e = Erosion2d(1, 1, 5, soft_max=False).cuda()
# prediction2 = e(prediction)
#
# precision2 = to_pil(prediction2.data.squeeze(0).cpu())
# precision2 = prediction2.data.squeeze(0).cpu().numpy()
# precision2 = precision2.resize(shape)
# prediction2 = np.array(precision2)
# prediction2 = prediction2.astype('float')
precision = to_pil(prediction.data.squeeze(0).cpu())
precision = precision.resize(shape)
prediction = np.array(precision)
prediction = prediction.astype('float')
# plt.style.use('classic')
# plt.subplot(1, 2, 1)
# plt.imshow(prediction)
# plt.subplot(1, 2, 2)
# plt.imshow(precision2[0])
# plt.show()
prediction = MaxMinNormalization(prediction, prediction.max(), prediction.min()) * 255.0
prediction = prediction.astype('uint8')
# if args['crf_refine']:
# prediction = crf_refine(np.array(img), prediction)
gt = np.array(Image.open(os.path.join(gt_root, img_name + '.png')).convert('L'))
precision, recall, mae = cal_precision_recall_mae(prediction, gt)
for pidx, pdata in enumerate(zip(precision, recall)):
p, r = pdata
precision_record[pidx].update(p)
recall_record[pidx].update(r)
mae_record.update(mae)
if args['save_results']:
folder, sub_name = os.path.split(img_name)
save_path = os.path.join(ckpt_path, exp_name, '(%s) %s_%s' % (exp_name, name, args['snapshot']), folder)
if not os.path.exists(save_path):
os.makedirs(save_path)
Image.fromarray(prediction).save(os.path.join(save_path, sub_name + '.png'))
fmeasure = cal_fmeasure([precord.avg for precord in precision_record],
[rrecord.avg for rrecord in recall_record])
results[name] = {'fmeasure': fmeasure, 'mae': mae_record.avg}
print ('test results:')
print (results)
log_path = os.path.join('result_all.txt')
open(log_path, 'a').write(exp_name + ' ' + args['snapshot'] + '\n')
open(log_path, 'a').write(str(results) + '\n\n')
def train(net, optimizer):
curr_iter = args['last_iter']
while True:
total_loss_record = AvgMeter()
lossL2H0_record, lossL2H1_record, lossL2H2_record, lossL2H3_record, lossL2H4_record = AvgMeter(
), AvgMeter(), AvgMeter(), AvgMeter(), AvgMeter()
lossH2L0_record, lossH2L1_record, lossH2L2_record, lossH2L3_record, lossH2L4_record = AvgMeter(
), AvgMeter(), AvgMeter(), AvgMeter(), AvgMeter()
for i, data in enumerate(train_loader):
optimizer.param_groups[0]['lr'] = 2 * args['lr'] * (
1 - float(curr_iter) / args['iter_num'])**args['lr_decay']
optimizer.param_groups[1]['lr'] = args['lr'] * (
1 - float(curr_iter) / args['iter_num'])**args['lr_decay']
inputs, labels = data
batch_size = inputs.size(0)
inputs = Variable(inputs).cuda()
labels = Variable(labels).cuda()
optimizer.zero_grad()
outputsL2H0, outputsL2H1, outputsL2H2, outputsL2H3, outputsL2H4, outputsH2L0, outputsH2L1, outputsH2L2, outputsH2L3, outputsH2L4, outputsFusion = net(
inputs)
lossL2H0 = criterion(outputsL2H0, labels)
lossL2H1 = criterion(outputsL2H1, labels)
lossL2H2 = criterion(outputsL2H2, labels)
lossL2H3 = criterion(outputsL2H3, labels)
lossL2H4 = criterion(outputsL2H4, labels)
lossH2L0 = criterion(outputsH2L0, labels)
lossH2L1 = criterion(outputsH2L1, labels)
lossH2L2 = criterion(outputsH2L2, labels)
lossH2L3 = criterion(outputsH2L3, labels)
lossH2L4 = criterion(outputsH2L4, labels)
lossFusion = criterion(outputsFusion, labels)
total_loss = lossFusion + lossL2H0 + lossL2H1 + lossL2H2 + lossL2H3 + lossL2H4 + lossH2L0 + lossH2L1 + lossH2L2 + lossH2L3 + lossH2L4
total_loss.backward()
optimizer.step()
total_loss_record.update(total_loss.data[0], batch_size)
lossL2H0_record.update(lossL2H0.data[0], batch_size)
lossL2H1_record.update(lossL2H1.data[0], batch_size)
lossL2H2_record.update(lossL2H2.data[0], batch_size)
lossL2H3_record.update(lossL2H3.data[0], batch_size)
lossL2H4_record.update(lossL2H4.data[0], batch_size)
lossH2L0_record.update(lossH2L0.data[0], batch_size)
lossH2L1_record.update(lossH2L1.data[0], batch_size)
lossH2L2_record.update(lossH2L2.data[0], batch_size)
lossH2L3_record.update(lossH2L3.data[0], batch_size)
lossH2L4_record.update(lossH2L4.data[0], batch_size)
curr_iter += 1
log = '[iter %d], [total loss %.5f], [lossL2H0 %.5f], [lossL2H1 %.5f], [lossL2H2 %.5f], [lossL2H3 %.5f], [lossL2H4 %.5f]' \
'[lossH2L0 %.5f], [lossH2L1 %.5f], [lossH2L2 %.5f], [lossH2L3 %.5f], [lossH2L4 %.5f], [lr %.13f]' % \
(curr_iter, total_loss_record.avg, lossL2H0_record.avg, lossL2H1_record.avg, lossL2H2_record.avg,
lossL2H3_record.avg, lossL2H4_record.avg, lossH2L0_record.avg, lossH2L1_record.avg, lossH2L2_record.avg,
lossH2L3_record.avg, lossH2L4_record.avg,
optimizer.param_groups[1]['lr'])
logWrite = '%d %.5f %.5f %.5f %.5f %.5f %.5f]' \
'%.5f %.5f %.5f %.5f %.5f %.13f]' % \
(curr_iter, total_loss_record.avg, lossL2H0_record.avg, lossL2H1_record.avg, lossL2H2_record.avg,
lossL2H3_record.avg, lossL2H4_record.avg, lossH2L0_record.avg, lossH2L1_record.avg, lossH2L2_record.avg,
lossH2L3_record.avg, lossH2L4_record.avg,
optimizer.param_groups[1]['lr'])
print log
open(log_path, 'a').write(logWrite + '\n')
if curr_iter == args['iter_num']:
torch.save(
net.state_dict(),
os.path.join(ckpt_path, exp_name, 'CA_%d.pth' % curr_iter))
torch.save(
optimizer.state_dict(),
os.path.join(ckpt_path, exp_name,
'CA_%d_optim.pth' % curr_iter))
return
def train(epoch):
print('\nEpoch: %d' % epoch)
net.train()
train_loss_record, bce_loss_record, dice_loss_record = AvgMeter(), AvgMeter(), AvgMeter()
for batch_idx, data in enumerate(train_loader):
if epoch == args['lr_step']:
optimizer.param_groups[0]['lr'] = 2 * args['lr'] / args['lr_decay']
optimizer.param_groups[1]['lr'] = args['lr'] / args['lr_decay']
inputs_volume, labels_volume = data
inputs_volume = norm_filter(inputs_volume) #normlization via means and std
sub_batch_len = math.ceil(inputs_volume.shape[1]/args['train_batch_size'])
print_flag=0
for sub_batch_idx in range(sub_batch_len):
#split volume to batchs
start = sub_batch_idx*args['train_batch_size']
end = (sub_batch_idx+1)*args['train_batch_size'] if (sub_batch_idx+1)*args['train_batch_size']<inputs_volume.shape[1] else inputs_volume.shape[1]
inputs = inputs_volume[:, start:end, :, :].permute(1, 0, 2, 3)
inputs = inputs.expand(torch.Size((inputs.shape[0], 3, inputs.shape[2], inputs.shape[3])))
labels = labels_volume[:, start:end, :, :].permute(1, 0, 2, 3)
if torch.max(labels).item() == 0:
print_flag += 1
print(print_flag)
continue
batch_size = inputs.size(0)
inputs = Variable(inputs).cuda()
labels = Variable(labels).cuda()
optimizer.zero_grad()
outputs = net(inputs)
# outputs = net(inputs)
# BCE loss and dice loss can be used
criterion_bce = nn.BCELoss()
criterion_dice = Dice_loss()
# if not isinstance(fnd_out, list):
loss_bce = criterion_bce(outputs, labels)
loss_dice = criterion_dice(outputs, labels)
# loss_bce = criterion_bce(outputs, labels)
# loss_dice = criterion_dice(outputs, labels)
# else:
# loss_bce = criterion_bce(outputs, labels)
# loss_dice = criterion_dice(outputs, labels)
# for fnd_mask, fpd_mask in zip(fnd_out, fpd_out):
# loss_bce += criterion_bce(fnd_mask, fnd) + criterion_bce(fpd_mask, fpd)
# else:
# loss_bce_each = [None] * len(outputs)
# loss_dice_each = [None] * len(outputs)
# for idx in range(len(outputs)):
# loss_bce_each[idx] = criterion_bce(outputs[idx], labels)
# loss_dice_each[idx] = criterion_dice(outputs[idx], labels)
# loss_bce = sum(loss_bce_each)
# loss_dice = sum(loss_dice_each)
# coeff = loss_dice.item()/(loss_bce.item()+1e-5) if loss_dice.item()/(loss_bce.item()+1e-5) < 1 else 1
coeff = 1
loss = coeff*loss_bce + loss_dice
# loss = loss_bce + loss_dice
loss.backward()
optimizer.step()
train_loss_record.update(loss.item(), batch_size)
bce_loss_record.update(loss_bce.item(), batch_size)
dice_loss_record.update(loss_dice.item(), batch_size)
log = 'iter: %d | [bce loss: %.5f], [dice loss: %.5f],[Total loss: %.5f], [lr: %.8f]' % \
(epoch, bce_loss_record.avg, dice_loss_record.avg, train_loss_record.avg, optimizer.param_groups[1]['lr'])
progress_bar(batch_idx, len(train_loader), log)
def train_online(net, seq_name='breakdance'):
online_args = {
'iter_num': 100,
'train_batch_size': 1,
'lr': 1e-10,
'lr_decay': 0.95,
'weight_decay': 5e-4,
'momentum': 0.95,
}
joint_transform = joint_transforms.Compose([
joint_transforms.ImageResize(380),
# joint_transforms.RandomCrop(473),
# joint_transforms.RandomHorizontallyFlip(),
# joint_transforms.RandomRotate(10)
])
target_transform = transforms.ToTensor()
# train_set = VideoFSImageFolder(to_test['davis'], seq_name, use_first=True, joint_transform=joint_transform, transform=img_transform)
train_set = VideoFirstImageFolder(to_test['davis'],
gt_root,
seq_name,
joint_transform=joint_transform,
transform=img_transform,
target_transform=target_transform)
online_train_loader = DataLoader(
train_set,
batch_size=online_args['train_batch_size'],
num_workers=1,
shuffle=False)
# criterion = nn.MSELoss().cuda()
criterion = nn.BCEWithLogitsLoss().cuda()
erosion = Erosion2d(1, 1, 5, soft_max=False).cuda()
net.train()
net.cuda()
# fix_parameters(net.named_parameters())
optimizer = optim.SGD([{
'params': [
param
for name, param in net.named_parameters() if name[-4:] == 'bias'
],
'lr':
2 * online_args['lr']
}, {
'params': [
param
for name, param in net.named_parameters() if name[-4:] != 'bias'
],
'lr':
online_args['lr'],
'weight_decay':
online_args['weight_decay']
}],
momentum=online_args['momentum'])
for curr_iter in range(0, online_args['iter_num']):
total_loss_record, loss0_record, loss1_record = AvgMeter(), AvgMeter(
), AvgMeter()
loss2_record = AvgMeter()
for i, data in enumerate(online_train_loader):
optimizer.param_groups[0]['lr'] = 2 * online_args['lr'] * (
1 - float(curr_iter) /
online_args['iter_num'])**online_args['lr_decay']
optimizer.param_groups[1]['lr'] = online_args['lr'] * (
1 - float(curr_iter) /
online_args['iter_num'])**online_args['lr_decay']
inputs, labels = data
batch_size = inputs.size(0)
inputs = Variable(inputs).cuda()
labels = Variable(labels).cuda()
optimizer.zero_grad()
if args['model'] == 'BASNet':
total_loss, loss0, loss1, loss2 = train_BASNet(
net, inputs, criterion, erosion, labels)
elif args['model'] == 'R3Net':
total_loss, loss0, loss1, loss2 = train_R3Net(
net, inputs, criterion, erosion, labels)
elif args['model'] == 'DSSNet':
total_loss, loss0, loss1, loss2 = train_DSSNet(
net, inputs, criterion, erosion, labels)
elif args['model'] == 'CPD':
total_loss, loss0, loss1, loss2 = train_CPD(
net, inputs, criterion, erosion, labels)
elif args['model'] == 'RAS':
total_loss, loss0, loss1, loss2 = train_RAS(
net, inputs, criterion, erosion, labels)
elif args['model'] == 'PoolNet':
total_loss, loss0, loss1, loss2 = train_PoolNet(
net, inputs, criterion, erosion, labels)
elif args['model'] == 'F3Net':
total_loss, loss0, loss1, loss2 = train_F3Net(
net, inputs, criterion, erosion, labels)
elif args['model'] == 'R2Net':
total_loss, loss0, loss1, loss2 = train_R2Net(
net, inputs, criterion, erosion, labels)
total_loss.backward()
optimizer.step()
total_loss_record.update(total_loss.data, batch_size)
loss0_record.update(loss0.data, batch_size)
loss1_record.update(loss1.data, batch_size)
loss2_record.update(loss2.data, batch_size)
# loss3_record.update(loss3.data, batch_size)
# loss4_record.update(loss4.data, batch_size)
log = '[iter %d], [total loss %.5f], [loss0 %.8f], [loss1 %.8f], [loss2 %.8f], [lr %.13f]' % \
(curr_iter, total_loss_record.avg, loss0_record.avg, loss1_record.avg, loss2_record.avg,
optimizer.param_groups[1]['lr'])
print(log)
print('taking snapshot ...')
torch.save(
net.state_dict(),
os.path.join(ckpt_path, exp_name,
str(args['snapshot']) + '_' + seq_name + '_online.pth'))
# torch.save(optimizer.state_dict(),
# os.path.join(ckpt_path, exp_name, '%d_optim.pth' % curr_iter))
return net
def train(net, optimizer):
curr_iter = args['last_iter']
while True:
total_loss_record, loss1_record, loss2_record, loss3_record, loss4_record, loss5_record, loss6_record, loss7_record, loss8_record = AvgMeter(
), AvgMeter(), AvgMeter(), AvgMeter(), AvgMeter(), AvgMeter(
), AvgMeter(), AvgMeter(), AvgMeter()
for i, data in enumerate(train_loader):
optimizer.param_groups[0]['lr'] = 2 * args['lr'] * (
1 - float(curr_iter) / args['iter_num'])**args['lr_decay']
optimizer.param_groups[1]['lr'] = args['lr'] * (
1 - float(curr_iter) / args['iter_num'])**args['lr_decay']
inputs, depth, labels = data
labels[labels > 0.5] = 1
labels[labels != 1] = 0
batch_size = inputs.size(0)
inputs = Variable(inputs).cuda()
depth = Variable(depth).cuda()
labels = Variable(labels).cuda()
outputs, outputs_fg, outputs_bg, attention1, attention2, attention3, attention4, attention5 = net(
inputs, depth) #hed
##########loss#############
optimizer.zero_grad()
labels1 = functional.interpolate(labels, size=24, mode='bilinear')
labels2 = functional.interpolate(labels, size=48, mode='bilinear')
labels3 = functional.interpolate(labels, size=96, mode='bilinear')
labels4 = functional.interpolate(labels, size=192, mode='bilinear')
loss1 = criterion_BCE(attention1, labels1)
loss2 = criterion_BCE(attention2, labels2)
loss3 = criterion_BCE(attention3, labels3)
loss4 = criterion_BCE(attention4, labels4)
loss5 = criterion_BCE(attention5, labels)
loss6 = criterion(outputs_fg, labels)
loss7 = criterion(outputs_bg, (1 - labels))
loss8 = criterion(outputs, labels)
total_loss = loss1 + loss2 + loss3 + loss4 + loss5 + loss6 + loss7 + loss8
total_loss.backward()
optimizer.step()
total_loss_record.update(total_loss.item(), batch_size)
loss1_record.update(loss1.item(), batch_size)
loss2_record.update(loss2.item(), batch_size)
loss3_record.update(loss3.item(), batch_size)
loss4_record.update(loss4.item(), batch_size)
loss5_record.update(loss5.item(), batch_size)
loss6_record.update(loss6.item(), batch_size)
loss7_record.update(loss7.item(), batch_size)
loss8_record.update(loss8.item(), batch_size)
curr_iter += 1
#############log###############
if curr_iter % 2050 == 0:
torch.save(
net.state_dict(),
os.path.join(ckpt_path, exp_name, '%d.pth' % curr_iter))
torch.save(
optimizer.state_dict(),
os.path.join(ckpt_path, exp_name,
'%d_optim.pth' % curr_iter))
log = '[iter %d], [total loss %.5f],[loss1 %.5f],,[loss2 %.5f],[loss3 %.5f],[loss4 %.5f],[loss5 %.5f],[loss6 %.5f],[loss7 %.5f],[loss8 %.5f],[lr %.13f] ' % \
(curr_iter, total_loss_record.avg, loss1_record.avg,loss2_record.avg,loss3_record.avg,loss4_record.avg,loss5_record.avg,loss6_record.avg,loss7_record.avg,loss8_record.avg,optimizer.param_groups[1]['lr'])
print(log)
open(log_path, 'a').write(log + '\n')
if curr_iter == args['iter_num']:
torch.save(
net.state_dict(),
os.path.join(ckpt_path, exp_name, '%d.pth' % curr_iter))
torch.save(
optimizer.state_dict(),
os.path.join(ckpt_path, exp_name,
'%d_optim.pth' % curr_iter))
return
import os
import numpy as np
from PIL import Image
from misc import check_mkdir, crf_refine, AvgMeter, cal_precision_recall_mae, cal_fmeasure
# root = '/home/qub/data/saliency/davis/davis_test2'
root_inference = '/home/ty/code/tf_saliency_attention/total_result/result_rnn_2018-08-04 11:08:00/'
root = '/home/ty/data/davis/480p/'
name = 'davis'
# gt_root = '/home/qub/data/saliency/davis/GT'
gt_root = '/home/ty/data/davis/GT/'
# gt_root = '/home/qub/data/saliency/VOS/GT'
precision_record, recall_record, = [AvgMeter() for _ in range(256)], [AvgMeter() for _ in range(256)]
mae_record = AvgMeter()
results = {}
# save_path = os.path.join(ckpt_path, exp_name, '(%s) %s_%s' % (exp_name, name, args['snapshot']))
folders = os.listdir(root_inference)
folders.sort()
for folder in folders:
imgs = os.listdir(os.path.join(root_inference, folder))
imgs.sort()
for img in imgs:
print(os.path.join(folder, img))
image = Image.open(os.path.join(root, folder, img[:-4] + '.jpg')).convert('RGB')
gt = np.array(Image.open(os.path.join(gt_root, folder, img)).convert('L'))
pred = np.array(Image.open(os.path.join(root_inference, folder, img)).convert('L'))
precision, recall, mae = cal_precision_recall_mae(pred, gt)
def train():
g = Generator(scale_factor=train_args['scale_factor']).cuda().train()
g = nn.DataParallel(g, device_ids=[0])
if len(train_args['g_snapshot']) > 0:
print 'load generator snapshot ' + train_args['g_snapshot']
g.load_state_dict(
torch.load(
os.path.join(train_args['ckpt_path'],
train_args['g_snapshot'])))
mse_criterion = nn.MSELoss().cuda()
g_mse_loss_record, psnr_record = AvgMeter(), AvgMeter()
iter_nums = len(train_loader)
if g_pretrain_args['pretrain']:
g_optimizer = optim.Adam(g.parameters(), lr=g_pretrain_args['lr'])
for epoch in range(g_pretrain_args['epoch_num']):
for i, data in enumerate(train_loader):
hr_imgs, _ = data
batch_size = hr_imgs.size(0)
lr_imgs = Variable(
torch.stack([train_lr_transform(img) for img in hr_imgs],
0)).cuda()
hr_imgs = Variable(hr_imgs).cuda()
g.zero_grad()
gen_hr_imgs = g(lr_imgs)
mse_loss = mse_criterion(gen_hr_imgs, hr_imgs)
mse_loss.backward()
g_optimizer.step()
g_mse_loss_record.update(mse_loss.data[0], batch_size)
psnr_record.update(10 * math.log10(1 / mse_loss.data[0]),
batch_size)
print '[pretrain]: [epoch %d], [iter %d / %d], [loss %.5f], [psnr %.5f]' % (
epoch + 1, i + 1, iter_nums, g_mse_loss_record.avg,
psnr_record.avg)
writer.add_scalar('pretrain_g_mse_loss', g_mse_loss_record.avg,
epoch * iter_nums + i + 1)
writer.add_scalar('pretrain_psnr', psnr_record.avg,
epoch * iter_nums + i + 1)
torch.save(
g.state_dict(),
os.path.join(
train_args['ckpt_path'],
'pretrain_g_epoch_%d_loss_%.5f_psnr_%.5f.pth' %
(epoch + 1, g_mse_loss_record.avg, psnr_record.avg)))
g_mse_loss_record.reset()
psnr_record.reset()
validate(g, epoch)
d = Discriminator().cuda().train()
d = nn.DataParallel(d, device_ids=[0])
if len(train_args['d_snapshot']) > 0:
print 'load discriminator snapshot ' + train_args['d_snapshot']
d.load_state_dict(
torch.load(
os.path.join(train_args['ckpt_path'],
train_args['d_snapshot'])))
g_optimizer = optim.RMSprop(g.parameters(), lr=train_args['g_lr'])
d_optimizer = optim.RMSprop(d.parameters(), lr=train_args['d_lr'])
perceptual_criterion, tv_criterion = PerceptualLoss().cuda(
), TotalVariationLoss().cuda()
g_mse_loss_record, g_perceptual_loss_record, g_tv_loss_record = AvgMeter(
), AvgMeter(), AvgMeter()
psnr_record, g_ad_loss_record, g_loss_record, d_loss_record = AvgMeter(
), AvgMeter(), AvgMeter(), AvgMeter()
for epoch in range(train_args['start_epoch'] - 1, train_args['epoch_num']):
for i, data in enumerate(train_loader):
hr_imgs, _ = data
batch_size = hr_imgs.size(0)
lr_imgs = Variable(
torch.stack([train_lr_transform(img) for img in hr_imgs],
0)).cuda()
hr_imgs = Variable(hr_imgs).cuda()
gen_hr_imgs = g(lr_imgs)
# update d
d.zero_grad()
d_ad_loss = d(gen_hr_imgs.detach()).mean() - d(hr_imgs).mean()
d_ad_loss.backward()
d_optimizer.step()
d_loss_record.update(d_ad_loss.data[0], batch_size)
for p in d.parameters():
p.data.clamp_(-train_args['c'], train_args['c'])
# update g
g.zero_grad()
g_mse_loss = mse_criterion(gen_hr_imgs, hr_imgs)
g_perceptual_loss = perceptual_criterion(gen_hr_imgs, hr_imgs)
g_tv_loss = tv_criterion(gen_hr_imgs)
g_ad_loss = -d(gen_hr_imgs).mean()
g_loss = g_mse_loss + 0.006 * g_perceptual_loss + 2e-8 * g_tv_loss + 0.001 * g_ad_loss
g_loss.backward()
g_optimizer.step()
g_mse_loss_record.update(g_mse_loss.data[0], batch_size)
g_perceptual_loss_record.update(g_perceptual_loss.data[0],
batch_size)
g_tv_loss_record.update(g_tv_loss.data[0], batch_size)
psnr_record.update(10 * math.log10(1 / g_mse_loss.data[0]),
batch_size)
g_ad_loss_record.update(g_ad_loss.data[0], batch_size)
g_loss_record.update(g_loss.data[0], batch_size)
print '[train]: [epoch %d], [iter %d / %d], [d_ad_loss %.5f], [g_ad_loss %.5f], [psnr %.5f], ' \
'[g_mse_loss %.5f], [g_perceptual_loss %.5f], [g_tv_loss %.5f] [g_loss %.5f]' % \
(epoch + 1, i + 1, iter_nums, d_loss_record.avg, g_ad_loss_record.avg, psnr_record.avg,
g_mse_loss_record.avg, g_perceptual_loss_record.avg, g_tv_loss_record.avg, g_loss_record.avg)
writer.add_scalar('d_loss', d_loss_record.avg,
epoch * iter_nums + i + 1)
writer.add_scalar('g_mse_loss', g_mse_loss_record.avg,
epoch * iter_nums + i + 1)
writer.add_scalar('g_perceptual_loss',
g_perceptual_loss_record.avg,
epoch * iter_nums + i + 1)
writer.add_scalar('g_tv_loss', g_tv_loss_record.avg,
epoch * iter_nums + i + 1)
writer.add_scalar('psnr', psnr_record.avg,
epoch * iter_nums + i + 1)
writer.add_scalar('g_ad_loss', g_ad_loss_record.avg,
epoch * iter_nums + i + 1)
writer.add_scalar('g_loss', g_loss_record.avg,
epoch * iter_nums + i + 1)
d_loss_record.reset()
g_mse_loss_record.reset()
g_perceptual_loss_record.reset()
g_tv_loss_record.reset()
psnr_record.reset()
g_ad_loss_record.reset()
g_loss_record.reset()
validate(g, epoch, d)
def train(net, optimizer):
curr_iter = args['last_iter']
while True:
train_loss_record = AvgMeter()
train_net_loss_record = AvgMeter()
train_depth_loss_record = AvgMeter()
for i, data in enumerate(train_loader):
optimizer.param_groups[0]['lr'] = 2 * args['lr'] * (
1 - float(curr_iter) / args['iter_num'])**args['lr_decay']
optimizer.param_groups[1]['lr'] = args['lr'] * (
1 - float(curr_iter) / args['iter_num'])**args['lr_decay']
inputs, gts, dps = data
batch_size = inputs.size(0)
inputs = Variable(inputs).cuda()
gts = Variable(gts).cuda()
dps = Variable(dps).cuda()
optimizer.zero_grad()
result, depth_pred = net(inputs)
loss_net = criterion(result, gts)
loss_depth = criterion_depth(depth_pred, dps)
loss = loss_net + loss_depth
loss.backward()
optimizer.step()
# for n, p in net.named_parameters():
# if n[-5:] == 'alpha':
# print(p.grad.data)
# print(p.data)
train_loss_record.update(loss.data, batch_size)
train_net_loss_record.update(loss_net.data, batch_size)
train_depth_loss_record.update(loss_depth.data, batch_size)
curr_iter += 1
log = '[iter %d], [train loss %.5f], [lr %.13f], [loss_net %.5f], [loss_depth %.5f]' % \
(curr_iter, train_loss_record.avg, optimizer.param_groups[1]['lr'],
train_net_loss_record.avg, train_depth_loss_record.avg)
print(log)
open(log_path, 'a').write(log + '\n')
if (curr_iter + 1) % args['val_freq'] == 0:
validate(net, curr_iter, optimizer)
if (curr_iter + 1) % args['snapshot_epochs'] == 0:
torch.save(
net.state_dict(),
os.path.join(ckpt_path, exp_name,
('%d.pth' % (curr_iter + 1))))
torch.save(
optimizer.state_dict(),
os.path.join(ckpt_path, exp_name,
('%d_optim.pth' % (curr_iter + 1))))
if curr_iter > args['iter_num']:
return
def train(net, optimizer):
start_time = time.time()
curr_iter = args['last_iter']
num_class = [0, 0, 0, 0, 0]
while True:
total_loss_record, loss0_record, loss1_record, loss2_record = AvgMeter(
), AvgMeter(), AvgMeter(), AvgMeter()
batch_time = AvgMeter()
end = time.time()
print('-----begining the first stage, train_mode==0-----')
for i, data in enumerate(train_loader):
optimizer.param_groups[0]['lr'] = 2 * args['lr'] * (
1 - float(curr_iter) / args['iter_num'])**args['lr_decay']
optimizer.param_groups[1]['lr'] = args['lr'] * (
1 - float(curr_iter) / args['iter_num'])**args['lr_decay']
inputs, gt, labels = data
print(labels)
# depends on the num of classes
cweight = torch.tensor([0.5, 0.75, 1, 1.25, 1.5])
#weight = torch.ones(size=gt.shape)
weight = gt.clone().detach()
sizec = labels.numpy()
#ta = np.zeros(shape=gt.shape)
'''
np.zeros(shape=labels.shape)
sc = gt.clone().detach()
for i in range(len(sizec)):
gta = np.array(to_pil(sc[i,:].data.squeeze(0).cpu()))#
#print(gta.shape)
labels[i] = cal_sc(gta)
sizec[i] = labels[i]
print(labels)
'''
batch_size = inputs.size(0)
inputs = Variable(inputs).cuda()
gt = Variable(gt).cuda()
labels = Variable(labels).cuda()
#print(sizec.shape)
optimizer.zero_grad()
p5, p4, p3, p2, p1, predict1, predict2, predict3, predict4, predict5, predict6, predict7, predict8, predict9, predict10, predict11 = net(
inputs, sizec) # mode=1
criterion = nn.BCEWithLogitsLoss().cuda()
criterion2 = nn.CrossEntropyLoss().cuda()
gt2 = gt.long()
gt2 = gt2.squeeze(1)
l5 = criterion2(p5, gt2)
l4 = criterion2(p4, gt2)
l3 = criterion2(p3, gt2)
l2 = criterion2(p2, gt2)
l1 = criterion2(p1, gt2)
loss0 = criterion(predict11, gt)
loss10 = criterion(predict10, gt)
loss9 = criterion(predict9, gt)
loss8 = criterion(predict8, gt)
loss7 = criterion(predict7, gt)
loss6 = criterion(predict6, gt)
loss5 = criterion(predict5, gt)
loss4 = criterion(predict4, gt)
loss3 = criterion(predict3, gt)
loss2 = criterion(predict2, gt)
loss1 = criterion(predict1, gt)
total_loss = l1 + l2 + l3 + l4 + l5 + loss0 + loss1 + loss2 + loss3 + loss4 + loss5 + loss6 + loss7 + loss8 + loss9 + loss10
total_loss.backward()
optimizer.step()
total_loss_record.update(total_loss.item(), batch_size)
loss1_record.update(l5.item(), batch_size)
loss0_record.update(loss0.item(), batch_size)
curr_iter += 1.0
batch_time.update(time.time() - end)
end = time.time()
log = '[iter %d], [R1/Mode0], [total loss %.5f]\n' \
'[l5 %.5f], [loss0 %.5f]\n' \
'[lr %.13f], [time %.4f]' % \
(curr_iter, total_loss_record.avg, loss1_record.avg, loss0_record.avg, optimizer.param_groups[1]['lr'],
batch_time.avg)
print(log)
print('Num of class:', num_class)
open(log_path, 'a').write(log + '\n')
if curr_iter == args['iter_num']:
torch.save(
net.state_dict(),
os.path.join(ckpt_path, exp_name, '%d.pth' % curr_iter))
torch.save(
optimizer.state_dict(),
os.path.join(ckpt_path, exp_name,
'%d_optim.pth' % curr_iter))
total_time = time.time() - start_time
print(total_time)
return
def inst_meter_dict(meter_list, meter_style='avg'):
result = dict()
for meter in meter_list:
if meter_style == 'avg':
result[meter] = AvgMeter()
return result
def main():
# net = R3Net(motion='', se_layer=False, dilation=False, basic_model='resnet50')
net = Distill(basic_model='resnet50', seq=True)
print ('load snapshot \'%s\' for testing' % args['snapshot'])
net.load_state_dict(torch.load(os.path.join(ckpt_path, exp_name, args['snapshot'] + '.pth'), map_location='cuda:2'))
net.eval()
net.cuda()
results = {}
with torch.no_grad():
for name, root in to_test.items():
precision_record, recall_record, = [AvgMeter() for _ in range(256)], [AvgMeter() for _ in range(256)]
mae_record = AvgMeter()
if args['save_results']:
check_mkdir(os.path.join(ckpt_path, exp_name, '(%s) %s_%s' % (exp_name, name, args['snapshot'])))
img_list = [i_id.strip() for i_id in open(imgs_path)]
video = ''
for idx, img_name in enumerate(img_list):
print ('predicting for %s: %d / %d' % (name, idx + 1, len(img_list)))
print(img_name)
# img_var = []
if args['seq']:
if video != img_name.split('/')[0]:
video = img_name.split('/')[0]
if name == 'VOS' or name == 'DAVSOD':
img = Image.open(os.path.join(root, img_name + '.png')).convert('RGB')
else:
img = Image.open(os.path.join(root, img_name + '.jpg')).convert('RGB')
shape = img.size
img = img.resize(args['input_size'])
img_var = Variable(img_transform(img).unsqueeze(0), volatile=True).cuda()
start = time.time()
_, _, prediction = net(img_var, img_var, flag='seq')
end = time.time()
print('running time:', (end - start))
else:
if name == 'VOS' or name == 'DAVSOD':
img = Image.open(os.path.join(root, img_name + '.png')).convert('RGB')
pre = Image.open(os.path.join(root, img_list[idx - 1] + '.png')).convert('RGB')
else:
img = Image.open(os.path.join(root, img_name + '.jpg')).convert('RGB')
pre = Image.open(os.path.join(root, img_list[idx - 1] + '.jpg')).convert('RGB')
shape = img.size
img = img.resize(args['input_size'])
pre = pre.resize(args['input_size'])
img_var = Variable(img_transform(img).unsqueeze(0), volatile=True).cuda()
pre_var = Variable(img_transform(pre).unsqueeze(0), volatile=True).cuda()
start = time.time()
_, _, prediction = net(pre_var, img_var, flag='seq')
end = time.time()
print('running time:', (end - start))
else:
start = time.time()
if name == 'VOS' or name == 'DAVSOD':
img = Image.open(os.path.join(root, img_name + '.png')).convert('RGB')
else:
img = Image.open(os.path.join(root, img_name + '.jpg')).convert('RGB')
shape = img.size
img = img.resize(args['input_size'])
img_var = Variable(img_transform(img).unsqueeze(0), volatile=True).cuda()
_, _, prediction, _ = net(img_var, img_var)
end = time.time()
print('running time:', (end - start))
precision = to_pil(prediction.data.squeeze(0).cpu())
precision = precision.resize(shape)
prediction = np.array(precision)
prediction = prediction.astype('float')
prediction = MaxMinNormalization(prediction, prediction.max(), prediction.min()) * 255.0
prediction = prediction.astype('uint8')
if args['crf_refine']:
prediction = crf_refine(np.array(img), prediction)
gt = np.array(Image.open(os.path.join(gt_root, img_name + '.png')).convert('L'))
precision, recall, mae = cal_precision_recall_mae(prediction, gt)
for pidx, pdata in enumerate(zip(precision, recall)):
p, r = pdata
precision_record[pidx].update(p)
recall_record[pidx].update(r)
mae_record.update(mae)
if args['save_results']:
folder, sub_name = os.path.split(img_name)
save_path = os.path.join(ckpt_path, exp_name, '(%s) %s_%s' % (exp_name, name, args['snapshot']), folder)
if not os.path.exists(save_path):
os.makedirs(save_path)
Image.fromarray(prediction).save(os.path.join(save_path, sub_name + '.png'))
fmeasure = cal_fmeasure([precord.avg for precord in precision_record],
[rrecord.avg for rrecord in recall_record])
results[name] = {'fmeasure': fmeasure, 'mae': mae_record.avg}
print ('test results:')
print (results)
def train(net, optimizer):
global total_epoch
curr_iter = 1
start_time = time.time()
for epoch in range(args['last_epoch'] + 1, args['last_epoch'] + 1 + args['epoch_num']):
loss_record, loss_4_record, loss_3_record, loss_2_record, loss_1_record, = AvgMeter(), AvgMeter(), AvgMeter(), AvgMeter(), AvgMeter()
train_iterator = tqdm(train_loader, total=len(train_loader))
for data in train_iterator:
if args['poly_train']:
base_lr = args['lr'] * (1 - float(curr_iter) / float(total_epoch)) ** args['lr_decay']
optimizer.param_groups[0]['lr'] = 2 * base_lr
optimizer.param_groups[1]['lr'] = 1 * base_lr
if args['poly_warmup']:
if curr_iter < args['warmup_epoch']:
base_lr = 1 / args['warmup_epoch'] * (1+curr_iter)
else:
curr_iter = curr_iter - args['warmup_epoch'] + 1
total_epoch = total_epoch - args['warmup_epoch'] + 1
base_lr = args['lr'] * (1 - float(curr_iter) / float(total_epoch)) ** args['lr_decay']
optimizer.param_groups[0]['lr'] = 2 * base_lr
optimizer.param_groups[1]['lr'] = 1 * base_lr
if args['cosine_warmup']:
if curr_iter < args['warmup_epoch']:
base_lr = 1 / args['warmup_epoch'] * (1+curr_iter)
else:
curr_iter = curr_iter - args['warmup_epoch'] + 1
total_epoch = total_epoch - args['warmup_epoch'] + 1
base_lr = args['lr'] * (1 + np.cos(np.pi *float(curr_iter) / float(total_epoch))) / 2
optimizer.param_groups[0]['lr'] = 2 * base_lr
optimizer.param_groups[1]['lr'] = 1 * base_lr
if args["f3_sche"]:
base_lr = args['lr'] *(1 - abs((curr_iter + 1) / (total_epoch + 1) * 2 - 1))
optimizer.param_groups[0]['lr'] = 2 * base_lr
optimizer.param_groups[1]['lr'] = 1 * base_lr
inputs, labels = data
batch_size = inputs.size(0)
inputs = Variable(inputs).cuda(device_ids[0])
labels = Variable(labels).cuda(device_ids[0])
optimizer.zero_grad()
predict_4, predict_3, predict_2, predict_1 = net(inputs)
loss_4 = bce_iou_edge_loss(predict_4, labels)
loss_3 = bce_iou_edge_loss(predict_3, labels)
loss_2 = bce_iou_edge_loss(predict_2, labels)
loss_1 = bce_iou_edge_loss(predict_1, labels)
loss = args['w2'][0] * loss_4 + args['w2'][1] * loss_3 + args['w2'][2] * loss_2 + args['w2'][3] * loss_1
loss.backward()
optimizer.step()
loss_record.update(loss.data, batch_size)
loss_4_record.update(loss_4.data, batch_size)
loss_3_record.update(loss_3.data, batch_size)
loss_2_record.update(loss_2.data, batch_size)
loss_1_record.update(loss_1.data, batch_size)
if curr_iter % 50 == 0:
writer.add_scalar('loss', loss, curr_iter)
writer.add_scalar('loss_4', loss_4, curr_iter)
writer.add_scalar('loss_3', loss_3, curr_iter)
writer.add_scalar('loss_2', loss_2, curr_iter)
writer.add_scalar('loss_1', loss_1, curr_iter)
log = '[%3d], [%6d], [%.6f], [%.5f], [%.5f], [%.5f], [%.5f], [%.5f]' % \
(epoch, curr_iter, base_lr, loss_record.avg, loss_4_record.avg, loss_3_record.avg,
loss_2_record.avg, loss_1_record.avg)
train_iterator.set_description(log)
open(log_path, 'a').write(log + '\n')
curr_iter += 1
if epoch in args['save_point']:
net.cpu()
torch.save(net.state_dict(), os.path.join(ckpt_path, exp_name, '%d.pth' % epoch))
net.cuda(device_ids[0])
if epoch >= args['epoch_num']:
net.cpu()
torch.save(net.state_dict(), os.path.join(ckpt_path, exp_name, '%d.pth' % epoch))
print("Total Training Time: {}".format(str(datetime.timedelta(seconds=int(time.time() - start_time)))))
print("Optimization Have Done!")
return
def train(net, optimizer):
curr_iter = 1
for epoch in range(args['last_epoch'] + 1,
args['last_epoch'] + 1 + args['epoch_num']):
loss_4_record, loss_3_record, loss_2_record, loss_1_record, \
loss_f_record, loss_record = AvgMeter(), AvgMeter(), AvgMeter(), AvgMeter(), AvgMeter(), AvgMeter()
train_iterator = tqdm(train_loader, total=len(train_loader))
for data in train_iterator:
if args['poly_train']:
base_lr = args['lr'] * (
1 - float(curr_iter) /
(args['epoch_num'] * len(train_loader)))**args['lr_decay']
optimizer.param_groups[0]['lr'] = 2 * base_lr
optimizer.param_groups[1]['lr'] = 1 * base_lr
inputs, labels = data
batch_size = inputs.size(0)
inputs = Variable(inputs).cuda(device_ids[0])
labels = Variable(labels).cuda(device_ids[0])
optimizer.zero_grad()
predict_4, predict_3, predict_2, predict_1, predict_f = net(inputs)
loss_4 = L.lovasz_hinge(predict_4, labels)
loss_3 = L.lovasz_hinge(predict_3, labels)
loss_2 = L.lovasz_hinge(predict_2, labels)
loss_1 = L.lovasz_hinge(predict_1, labels)
loss_f = L.lovasz_hinge(predict_f, labels)
loss = loss_4 + loss_3 + loss_2 + loss_1 + loss_f
loss.backward()
optimizer.step()
loss_record.update(loss.data, batch_size)
loss_4_record.update(loss_4.data, batch_size)
loss_3_record.update(loss_3.data, batch_size)
loss_2_record.update(loss_2.data, batch_size)
loss_1_record.update(loss_1.data, batch_size)
loss_f_record.update(loss_f.data, batch_size)
if curr_iter % 50 == 0:
writer.add_scalar('loss', loss, curr_iter)
writer.add_scalar('loss_4', loss_4, curr_iter)
writer.add_scalar('loss_3', loss_3, curr_iter)
writer.add_scalar('loss_2', loss_2, curr_iter)
writer.add_scalar('loss_1', loss_1, curr_iter)
writer.add_scalar('loss_f', loss_f, curr_iter)
log = '[%3d], [%6d], [%.6f], [%.5f], [L4: %.5f], [L3: %.5f], [L2: %.5f], [L1: %.5f], [Lf: %.5f]' % \
(epoch, curr_iter, base_lr, loss_record.avg, loss_4_record.avg, loss_3_record.avg, loss_2_record.avg,
loss_1_record.avg, loss_f_record.avg)
train_iterator.set_description(log)
open(log_path, 'a').write(log + '\n')
curr_iter += 1
if epoch in args['save_point']:
net.cpu()
torch.save(net.module.state_dict(),
os.path.join(ckpt_path, exp_name, '%d.pth' % epoch))
net.cuda(device_ids[0])
if epoch >= args['epoch_num']:
net.cpu()
torch.save(net.module.state_dict(),
os.path.join(ckpt_path, exp_name, '%d.pth' % epoch))
print("Optimization Have Done!")
return
def validate(g, curr_epoch, d=None):
g.eval()
mse_criterion = nn.MSELoss()
g_mse_loss_record, psnr_record = AvgMeter(), AvgMeter()
for name, loader in val_loader.iteritems():
val_visual = []
# note that the batch size is 1
for i, data in enumerate(loader):
hr_img, _ = data
lr_img, hr_restore_img = val_lr_transform(hr_img.squeeze(0))
lr_img = Variable(lr_img.unsqueeze(0), volatile=True).cuda()
hr_restore_img = hr_restore_img
hr_img = Variable(hr_img, volatile=True).cuda()
gen_hr_img = g(lr_img)
g_mse_loss = mse_criterion(gen_hr_img, hr_img)
g_mse_loss_record.update(g_mse_loss.data[0])
psnr_record.update(10 * math.log10(1 / g_mse_loss.data[0]))
val_visual.extend([
val_display_transform(hr_restore_img),
val_display_transform(hr_img.cpu().data.squeeze(0)),
val_display_transform(gen_hr_img.cpu().data.squeeze(0))
])
val_visual = torch.stack(val_visual, 0)
val_visual = vutils.make_grid(val_visual, nrow=3, padding=5)
snapshot_name = 'epoch_%d_%s_g_mse_loss_%.5f_psnr_%.5f' % (
curr_epoch + 1, name, g_mse_loss_record.avg, psnr_record.avg)
if d is None:
snapshot_name = 'pretrain_' + snapshot_name
writer.add_scalar('pretrain_validate_%s_psnr' % name,
psnr_record.avg, curr_epoch + 1)
writer.add_scalar('pretrain_validate_%s_g_mse_loss' % name,
g_mse_loss_record.avg, curr_epoch + 1)
print '[pretrain validate %s]: [epoch %d], [g_mse_loss %.5f], [psnr %.5f]' % (
name, curr_epoch + 1, g_mse_loss_record.avg, psnr_record.avg)
else:
writer.add_scalar('validate_%s_psnr' % name, psnr_record.avg,
curr_epoch + 1)
writer.add_scalar('validate_%s_g_mse_loss' % name,
g_mse_loss_record.avg, curr_epoch + 1)
print '[validate %s]: [epoch %d], [g_mse_loss %.5f], [psnr %.5f]' % (
name, curr_epoch + 1, g_mse_loss_record.avg, psnr_record.avg)
torch.save(
d.state_dict(),
os.path.join(train_args['ckpt_path'],
snapshot_name + '_d.pth'))
torch.save(
g.state_dict(),
os.path.join(train_args['ckpt_path'], snapshot_name + '_g.pth'))
writer.add_image(snapshot_name, val_visual)
g_mse_loss_record.reset()
psnr_record.reset()
g.train()
def train(net, optimizer):
curr_iter = args['last_iter']
train_loss_record = AvgMeter()
train_net_loss_record = AvgMeter()
train_uncertainty_loss_record = AvgMeter()
while True:
for i, data in enumerate(train_loader):
optimizer.param_groups[0]['lr'] = 2 * args['lr'] * (
1 - float(curr_iter) / args['iter_num'])**args['lr_decay']
optimizer.param_groups[1]['lr'] = args['lr'] * (
1 - float(curr_iter) / args['iter_num'])**args['lr_decay']
inputs, gts = data
batch_size = inputs.size(0)
inputs = Variable(inputs).cuda()
gts = Variable(gts).cuda()
optimizer.zero_grad()
result = net(inputs)
nhw = inputs.size(0) * inputs.size(2) * inputs.size(3)
loss_net = criterion(result, gts)
loss = loss_net
loss.backward()
optimizer.step()
#print(uncertainty)
train_loss_record.update(loss.data, batch_size)
train_net_loss_record.update(loss_net.data, batch_size)
#train_uncertainty_loss_record.update(loss_uncertainty_regular.data, batch_size)
curr_iter += 1
# log = '[iter %d], [train loss %.5f], [lr %.8f], [loss_net %.5f], [w1 %.8f], [loss_uncertainty %.5f]' % \
# (curr_iter, train_loss_record.avg, optimizer.param_groups[1]['lr'],
# train_net_loss_record.avg, w1, train_uncertainty_loss_record.avg)
log = '[iter %d], [train loss %.5f], [lr %.8f], [loss_net %.5f]' % \
(curr_iter, train_loss_record.avg, optimizer.param_groups[1]['lr'],
train_net_loss_record.avg)
print(log)
open(log_path, 'a').write(log + '\n')
if (curr_iter + 1) % args['val_freq'] == 0:
validate(net, curr_iter, optimizer)
if (curr_iter + 1) % args['snapshot_epochs'] == 0:
torch.save(
net.state_dict(),
os.path.join(ckpt_path, exp_name,
('%d.pth' % (curr_iter + 1))))
torch.save(
optimizer.state_dict(),
os.path.join(ckpt_path, exp_name,
('%d_optim.pth' % (curr_iter + 1))))
if curr_iter > args['iter_num']:
return
def train(net, optimizer):
curr_iter = 1
for epoch in range(args['last_epoch'] + 1,
args['last_epoch'] + 1 + args['epoch_num']):
loss_record, loss_b_record, loss_c_record, loss_o_record = AvgMeter(
), AvgMeter(), AvgMeter(), AvgMeter()
train_iterator = tqdm(train_loader, total=len(train_loader))
for data in train_iterator:
if args['poly_train']:
base_lr = args['lr'] * (
1 - float(curr_iter) /
(args['epoch_num'] * len(train_loader)))**args['lr_decay']
optimizer.param_groups[0]['lr'] = 2 * base_lr
optimizer.param_groups[1]['lr'] = 1 * base_lr
inputs, labels, edges = data
batch_size = inputs.size(0)
inputs = Variable(inputs).cuda(device_ids[0])
labels = Variable(labels).cuda(device_ids[0])
edges = Variable(edges).cuda(device_ids[0])
optimizer.zero_grad()
predict_c, predict_b, predict_o = net(inputs)
loss_b = bce(predict_b, edges)
loss_c = L.lovasz_hinge(predict_c, labels)
loss_o = L.lovasz_hinge(predict_o, labels)
loss = loss_b + loss_c + loss_o
loss.backward()
optimizer.step()
loss_record.update(loss.data, batch_size)
loss_b_record.update(loss_b.data, batch_size)
loss_c_record.update(loss_c.data, batch_size)
loss_o_record.update(loss_o.data, batch_size)
if curr_iter % 50 == 0:
writer.add_scalar('loss', loss, curr_iter)
writer.add_scalar('loss_b', loss_b, curr_iter)
writer.add_scalar('loss_c', loss_c, curr_iter)
writer.add_scalar('loss_o', loss_o, curr_iter)
log = '[Epoch: %2d], [Iter: %5d], [%.7f], [Sum: %.5f], [Lb: %.5f], [Lc: %.5f], [Lo: %.5f]' % \
(epoch, curr_iter, base_lr, loss_record.avg, loss_b_record.avg, loss_c_record.avg, loss_o_record.avg)
train_iterator.set_description(log)
open(log_path, 'a').write(log + '\n')
curr_iter += 1
if epoch in args['save_point']:
net.cpu()
torch.save(net.module.state_dict(),
os.path.join(ckpt_path, exp_name, '%d.pth' % epoch))
net.cuda(device_ids[0])
if epoch >= args['epoch_num']:
net.cpu()
torch.save(net.module.state_dict(),
os.path.join(ckpt_path, exp_name, '%d.pth' % epoch))
print("Optimization Have Done!")
return
train_set = VideoImageFolder(video_train_path, imgs_file, joint_transform, img_transform, target_transform)
else:
train_set = VideoImage2Folder(video_train_path, imgs_file, video_seq_path + '/DAFB2', video_seq_gt_path + '/DAFB2',
joint_transform, None, input_size, img_transform, target_transform)
train_loader = DataLoader(train_set, batch_size=args['train_batch_size'], num_workers=4, shuffle=True)
criterion = nn.BCEWithLogitsLoss().cuda()
if args['L2']:
criterion_l2 = nn.MSELoss().cuda()
# criterion_pair = CriterionPairWise(scale=0.5).cuda()
if args['KL']:
criterion_kl = CriterionKL3().cuda()
log_path = os.path.join(ckpt_path, exp_name, str(datetime.datetime.now()) + '.txt')
total_loss_record, loss0_record, loss1_record, loss2_record = AvgMeter(), AvgMeter(), AvgMeter(), AvgMeter()
def fix_parameters(parameters):
for name, parameter in parameters:
if name.find('motion') >= 0 \
or name.find('GRU') >= 0 or name.find('predict') >= 0:
print(name, 'is not fixed')
else:
print(name, 'is fixed')
parameter.requires_grad = False
def main():
net = Distill(basic_model=args['basic_model'], seq=True).cuda().train()
def main():
net = ConcatNet(backbone='resnet50', embedding=128, batch_mode='old')
print('load snapshot \'%s\' for testing' % args['snapshot'])
net.load_state_dict(
torch.load(os.path.join(ckpt_path, exp_name,
args['snapshot'] + '.pth'),
map_location='cuda:0'))
net.eval()
net.cuda()
results = {}
with torch.no_grad():
precision_record, recall_record, = [AvgMeter() for _ in range(256)], [
AvgMeter() for _ in range(256)
]
mae_record = AvgMeter()
if args['save_results']:
check_mkdir(
os.path.join(
ckpt_path, exp_name,
'(%s) %s_%s' % (exp_name, 'DAVIS', args['snapshot'])))
index = 0
old_seq = ''
# img_list = [os.path.splitext(f)[0] for f in os.listdir(root) if f.endswith('.jpg')]
for idx, batch in enumerate(testloader):
print('%d processd ' % (idx))
target = batch['target']
seq_name = batch['seq_name']
args['seq_name'] = seq_name[0]
print('sequence name:', seq_name[0])
if old_seq == args['seq_name']:
index = index + 1
else:
index = 0
output_sum = 0
for i in range(0, args['sample_range']):
search = batch['search_' + str(i)]
output = net(
Variable(target, volatile=True).cuda(),
Variable(search, volatile=True).cuda())
output = F.upsample(output,
size=target.size()[2:],
mode='bilinear',
align_corners=True)
output = F.sigmoid(output)
output_sum = output_sum + output[0].data.cpu().numpy()
output_final = output_sum / args['sample_range']
output_final = cv2.resize(output_final[0], (854, 480))
output_final = (output_final * 255).astype(np.uint8)
gt = np.array(
Image.open(
os.path.join(args['gt_dir'], args['seq_name'],
str(index).zfill(5) + '.png')).convert('L'))
precision, recall, mae = cal_precision_recall_mae(output_final, gt)
for pidx, pdata in enumerate(zip(precision, recall)):
p, r = pdata
precision_record[pidx].update(p)
recall_record[pidx].update(r)
mae_record.update(mae)
if args['save_results']:
old_seq = args['seq_name']
save_path = os.path.join(
ckpt_path, exp_name,
'(%s) %s_%s' % (exp_name, 'DAVIS', args['snapshot']),
args['seq_name'])
if not os.path.exists(save_path):
os.makedirs(save_path)
Image.fromarray(output_final).save(
os.path.join(save_path,
str(index).zfill(5) + '.png'))
fmeasure = cal_fmeasure([precord.avg for precord in precision_record],
[rrecord.avg for rrecord in recall_record])
results['DAVIS'] = {'fmeasure': fmeasure, 'mae': mae_record.avg}
print('test results:')
print(results)
def train(net, optimizer):
curr_iter = args['last_iter']
while True:
train_loss_record, loss_fuse_record, loss1_h2l_record = AvgMeter(
), AvgMeter(), AvgMeter()
loss2_h2l_record, loss3_h2l_record, loss4_h2l_record = AvgMeter(
), AvgMeter(), AvgMeter()
loss1_l2h_record, loss2_l2h_record, loss3_l2h_record = AvgMeter(
), AvgMeter(), AvgMeter()
loss4_l2h_record = AvgMeter()
for i, data in enumerate(train_loader):
optimizer.param_groups[0]['lr'] = 2 * args['lr'] * (
1 - float(curr_iter) / args['iter_num'])**args['lr_decay']
optimizer.param_groups[1]['lr'] = args['lr'] * (
1 - float(curr_iter) / args['iter_num'])**args['lr_decay']
inputs, labels = data
batch_size = inputs.size(0)
inputs = Variable(inputs).cuda()
labels = Variable(labels).cuda()
optimizer.zero_grad()
fuse_predict, predict1_h2l, predict2_h2l, predict3_h2l, predict4_h2l, \
predict1_l2h, predict2_l2h, predict3_l2h, predict4_l2h = net(inputs)
loss_fuse = bce_logit(fuse_predict, labels)
loss1_h2l = bce_logit(predict1_h2l, labels)
loss2_h2l = bce_logit(predict2_h2l, labels)
loss3_h2l = bce_logit(predict3_h2l, labels)
loss4_h2l = bce_logit(predict4_h2l, labels)
loss1_l2h = bce_logit(predict1_l2h, labels)
loss2_l2h = bce_logit(predict2_l2h, labels)
loss3_l2h = bce_logit(predict3_l2h, labels)
loss4_l2h = bce_logit(predict4_l2h, labels)
loss = loss_fuse + loss1_h2l + loss2_h2l + loss3_h2l + loss4_h2l + loss1_l2h + \
loss2_l2h + loss3_l2h + loss4_l2h
loss.backward()
optimizer.step()
train_loss_record.update(loss.data, batch_size)
loss_fuse_record.update(loss_fuse.data, batch_size)
loss1_h2l_record.update(loss1_h2l.data, batch_size)
loss2_h2l_record.update(loss2_h2l.data, batch_size)
loss3_h2l_record.update(loss3_h2l.data, batch_size)
loss4_h2l_record.update(loss4_h2l.data, batch_size)
loss1_l2h_record.update(loss1_l2h.data, batch_size)
loss2_l2h_record.update(loss2_l2h.data, batch_size)
loss3_l2h_record.update(loss3_l2h.data, batch_size)
loss4_l2h_record.update(loss4_l2h.data, batch_size)
curr_iter += 1
log = '[iter %d], [train loss %.5f], [loss_fuse %.5f], [loss1_h2l %.5f], [loss2_h2l %.5f], ' \
'[loss3_h2l %.5f], [loss4_h2l %.5f], [loss1_l2h %.5f], [loss2_l2h %.5f], [loss3_l2h %.5f], ' \
'[loss4_l2h %.5f], [lr %.13f]' % \
(curr_iter, train_loss_record.avg, loss_fuse_record.avg, loss1_h2l_record.avg, loss2_h2l_record.avg,
loss3_h2l_record.avg, loss4_h2l_record.avg, loss1_l2h_record.avg, loss2_l2h_record.avg,
loss3_l2h_record.avg, loss4_l2h_record.avg, optimizer.param_groups[1]['lr'])
print log
open(log_path, 'a').write(log + '\n')
if curr_iter > args['iter_num']:
torch.save(
net.state_dict(),
os.path.join(ckpt_path, exp_name, '%d.pth' % curr_iter))
return
def train(exp_name):
log_path = os.path.join(ckpt_path, exp_name,
str(datetime.datetime.now()) + '.txt')
net = DPNet().cuda().train()
optimizer = optim.SGD([{
'params': [
param
for name, param in net.named_parameters() if name[-4:] == 'bias'
],
'lr':
2 * args['lr']
}, {
'params': [
param
for name, param in net.named_parameters() if name[-4:] != 'bias'
],
'lr':
args['lr'],
'weight_decay':
args['weight_decay']
}],
momentum=args['momentum'])
if len(args['snapshot']) > 0:
print('training resumes from ' + args['snapshot'])
net.load_state_dict(
torch.load(
os.path.join(ckpt_path, exp_name, args['snapshot'] + '.pth')))
optimizer.param_groups[0]['lr'] = 2 * args['lr']
optimizer.param_groups[1]['lr'] = args['lr']
check_mkdir(ckpt_path)
check_mkdir(os.path.join(ckpt_path, exp_name))
open(log_path, 'w').write(str(args) + '\n\n')
print 'start to train'
curr_iter = args['last_iter']
while True:
total_loss_record = AvgMeter()
loss1_record, loss2_record, loss3_record, loss4_record = AvgMeter(
), AvgMeter(), AvgMeter(), AvgMeter()
loss_DPM1_record = AvgMeter()
for i, data in enumerate(train_loader):
optimizer.param_groups[0]['lr'] = 2 * args['lr'] * (
1 - float(curr_iter) / args['iter_num'])**args['lr_decay']
optimizer.param_groups[1]['lr'] = args['lr'] * (
1 - float(curr_iter) / args['iter_num'])**args['lr_decay']
inputs, labels = data
batch_size = inputs.size(0)
inputs = Variable(inputs).cuda()
labels = Variable(labels).cuda()
optimizer.zero_grad()
predict1, predict2, predict3, predict4, predict_DPM1 = net(inputs)
loss1 = criterionBCE(predict1, labels)
loss2 = criterionBCE(predict2, labels)
loss3 = criterionBCE(predict3, labels)
loss4 = criterionBCE(predict4, labels)
loss_DPM1 = criterionBCE(predict_DPM1, labels)
total_loss = loss1 + loss2 + loss3 + loss4 + loss_DPM1
total_loss.backward()
optimizer.step()
total_loss_record.update(total_loss.item(), batch_size)
loss1_record.update(loss1.item(), batch_size)
loss2_record.update(loss2.item(), batch_size)
loss3_record.update(loss3.item(), batch_size)
loss4_record.update(loss4.item(), batch_size)
loss_DPM1_record.update(loss_DPM1.item(), batch_size)
curr_iter += 1
log = '[iter %d], [total loss %.5f], [loss1 %.5f], [loss2 %.5f], [loss3 %.5f], ' \
'[loss4 %.5f], [loss_DPM1 %.5f], [lr %.13f]' \
% (curr_iter, total_loss_record.avg, loss1_record.avg, loss2_record.avg, loss3_record.avg,
loss4_record.avg, loss_DPM1_record.avg, optimizer.param_groups[1]['lr'])
print log
open(log_path, 'a').write(log + '\n')
if curr_iter == args['iter_num']:
torch.save(
net.state_dict(),
os.path.join(ckpt_path, exp_name, '%d.pth' % curr_iter))
return
def main():
if args['model'] == 'BASNet':
net = BASNet(3, 1)
elif args['model'] == 'R3Net':
net = R3Net()
elif args['model'] == 'DSSNet':
net = build_model()
elif args['model'] == 'CPD':
net = CPD_ResNet()
elif args['model'] == 'RAS':
net = RAS()
elif args['model'] == 'PiCANet':
net = Unet()
elif args['model'] == 'PoolNet':
net = build_model_poolnet(base_model_cfg='resnet')
elif args['model'] == 'R2Net':
net = build_model_r2net(base_model_cfg='resnet')
elif args['model'] == 'F3Net':
net = F3Net(cfg=None)
print('load snapshot \'%s\' for testing' % args['snapshot'])
net.load_state_dict(
torch.load(os.path.join(ckpt_path, exp_name,
args['snapshot'] + '.pth'),
map_location='cuda:0'))
# net = train_online(net)
results = {}
for name, root in to_test.items():
precision_record, recall_record, = [AvgMeter() for _ in range(256)], [
AvgMeter() for _ in range(256)
]
mae_record = AvgMeter()
if args['save_results']:
check_mkdir(
os.path.join(
ckpt_path, exp_name, '(%s) %s_%s' %
(exp_name, name, args['snapshot'] + '_online')))
folders = os.listdir(root)
folders.sort()
folders = ['bmx-trees']
for folder in folders:
if args['online_train']:
net = train_online(net, seq_name=folder)
net.load_state_dict(
torch.load(os.path.join(
ckpt_path, exp_name,
str(args['snapshot']) + '_' + folder + '_online.pth'),
map_location='cuda:0'))
with torch.no_grad():
net.eval()
net.cuda()
imgs = os.listdir(os.path.join(root, folder))
imgs.sort()
for i in range(args['start'], len(imgs)):
print(imgs[i])
start = time.time()
img = Image.open(os.path.join(root, folder,
imgs[i])).convert('RGB')
shape = img.size
img = img.resize(args['input_size'])
img_var = Variable(img_transform(img).unsqueeze(0),
volatile=True).cuda()
if args['model'] == 'BASNet':
prediction, _, _, _, _, _, _, _ = net(img_var)
prediction = torch.sigmoid(prediction)
elif args['model'] == 'R3Net':
prediction = net(img_var)
elif args['model'] == 'DSSNet':
select = [1, 2, 3, 6]
prediction = net(img_var)
prediction = torch.mean(torch.cat(
[torch.sigmoid(prediction[i]) for i in select],
dim=1),
dim=1,
keepdim=True)
elif args['model'] == 'CPD':
_, prediction = net(img_var)
prediction = torch.sigmoid(prediction)
elif args['model'] == 'RAS':
prediction, _, _, _, _ = net(img_var)
prediction = torch.sigmoid(prediction)
elif args['model'] == 'PoolNet':
prediction = net(img_var)
prediction = torch.sigmoid(prediction)
elif args['model'] == 'F3Net':
_, prediction, _, _, _, _ = net(img_var)
prediction = torch.sigmoid(prediction)
elif args['model'] == 'R2Net':
_, _, _, _, _, prediction = net(img_var)
prediction = torch.sigmoid(prediction)
end = time.time()
print('running time:', (end - start))
if args['crf_refine']:
prediction = crf_refine(np.array(img), prediction)
precision = to_pil(prediction.data.squeeze(0).cpu())
precision = precision.resize(shape)
prediction = np.array(precision)
prediction = prediction.astype('float')
prediction = MaxMinNormalization(
prediction, prediction.max(), prediction.min()) * 255.0
prediction = prediction.astype('uint8')
gt = np.array(
Image.open(
os.path.join(gt_root, folder,
imgs[i][:-4] + '.png')).convert('L'))
precision, recall, mae = cal_precision_recall_mae(
prediction, gt)
for pidx, pdata in enumerate(zip(precision, recall)):
p, r = pdata
precision_record[pidx].update(p)
recall_record[pidx].update(r)
mae_record.update(mae)
if args['save_results']:
# folder, sub_name = os.path.split(imgs[i])
save_path = os.path.join(
ckpt_path, exp_name, '(%s) %s_%s' %
(exp_name, name, args['snapshot'] + '_online'),
folder)
if not os.path.exists(save_path):
os.makedirs(save_path)
Image.fromarray(prediction).save(
os.path.join(save_path, imgs[i]))
fmeasure = cal_fmeasure([precord.avg for precord in precision_record],
[rrecord.avg for rrecord in recall_record])
results[name] = {'fmeasure': fmeasure, 'mae': mae_record.avg}
print('test results:')
print(results)
def main():
net = R3Net_prior(motion='GRU',
se_layer=False,
attention=False,
pre_attention=True,
basic_model='resnet50',
sta=False,
naive_fuse=False)
print('load snapshot \'%s\' for testing' % args['snapshot'])
net.load_state_dict(
torch.load(os.path.join(ckpt_path, exp_name,
args['snapshot'] + '.pth'),
map_location='cuda:1'))
net.eval()
net.cuda()
results = {}
with torch.no_grad():
for name, root in to_test.items():
precision_record, recall_record, = [
AvgMeter() for _ in range(256)
], [AvgMeter() for _ in range(256)]
mae_record = AvgMeter()
if args['save_results']:
check_mkdir(
os.path.join(
ckpt_path, exp_name,
'(%s) %s_%s' % (exp_name, name, args['snapshot'])))
img_list = [i_id.strip() for i_id in open(imgs_path)]
# img_list = [os.path.splitext(f)[0] for f in os.listdir(root) if f.endswith('.jpg')]
for idx, img_names in enumerate(img_list):
print('predicting for %s: %d / %d' %
(name, idx + 1, len(img_list)))
img_seq = img_names.split(',')
for ratio in args['scale_ratio']:
prediction_scale = []
img_var = []
for img_name in img_seq:
if name == 'VOS' or name == 'DAVSOD':
img = Image.open(
os.path.join(root,
img_name + '.png')).convert('RGB')
else:
img = Image.open(
os.path.join(root,
img_name + '.jpg')).convert('RGB')
shape = img.size
new_dims = (int(img.size[0] * ratio),
int(img.size[1] * ratio))
img = img.resize(new_dims, Image.BILINEAR)
img = img.resize(args['input_size'])
img_var.append(
Variable(img_transform(img).unsqueeze(0),
volatile=True).cuda())
img_var = torch.cat(img_var, dim=0)
start = time.time()
prediction = net(img_var)
end = time.time()
print('running time:', (end - start))
prediction_scale.append(prediction)
prediction = torch.mean(torch.stack(prediction_scale, dim=0),
0)
prediction = to_pil(prediction.data.squeeze(0).cpu())
prediction = prediction.resize(shape)
prediction = np.array(prediction)
prediction = prediction.astype('float')
prediction = MaxMinNormalization(prediction, prediction.max(),
prediction.min()) * 255.0
prediction = prediction.astype('uint8')
if args['crf_refine']:
prediction = crf_refine(np.array(img), prediction)
gt = np.array(
Image.open(os.path.join(gt_root, img_seq[-1] +
'.png')).convert('L'))
precision, recall, mae = cal_precision_recall_mae(
prediction, gt)
for pidx, pdata in enumerate(zip(precision, recall)):
p, r = pdata
precision_record[pidx].update(p)
recall_record[pidx].update(r)
mae_record.update(mae)
if args['save_results']:
folder, sub_name = os.path.split(img_name)
save_path = os.path.join(
ckpt_path, exp_name,
'(%s) %s_%s' % (exp_name, name, args['snapshot']),
folder)
if not os.path.exists(save_path):
os.makedirs(save_path)
Image.fromarray(prediction).save(
os.path.join(save_path, sub_name + '.png'))
fmeasure = cal_fmeasure(
[precord.avg for precord in precision_record],
[rrecord.avg for rrecord in recall_record])
results[name] = {'fmeasure': fmeasure, 'mae': mae_record.avg}
print('test results:')
print(results)
def main():
net = AADFNet().cuda()
net = nn.DataParallel(net, device_ids=[0])
print exp_name + 'crf: ' + str(args['crf_refine'])
print 'load snapshot \'%s\' for testing' % args['snapshot']
net.load_state_dict(
torch.load(os.path.join(ckpt_path, exp_name,
args['snapshot'] + '.pth')))
net.eval()
with torch.no_grad():
results = {}
for name, root in to_test.iteritems():
precision_record, recall_record, = [
AvgMeter() for _ in range(256)
], [AvgMeter() for _ in range(256)]
mae_record = AvgMeter()
time_record = AvgMeter()
img_list = [
os.path.splitext(f)[0] for f in os.listdir(root)
if f.endswith('.jpg')
]
for idx, img_name in enumerate(img_list):
img_name = img_list[idx]
print 'predicting for %s: %d / %d' % (name, idx + 1,
len(img_list))
check_mkdir(
os.path.join(
ckpt_path, exp_name,
'(%s) %s_%s' % (exp_name, name, args['snapshot'])))
start = time.time()
img = Image.open(os.path.join(root, img_name +
'.jpg')).convert('RGB')
img_var = Variable(img_transform(img).unsqueeze(0),
volatile=True).cuda()
prediction = net(img_var)
W, H = img.size
prediction = F.upsample_bilinear(prediction, size=(H, W))
prediction = np.array(to_pil(prediction.data.squeeze(0).cpu()))
if args['crf_refine']:
prediction = crf_refine(np.array(img), prediction)
end = time.time()
gt = np.array(
Image.open(os.path.join(root,
img_name + '.png')).convert('L'))
precision, recall, mae = cal_precision_recall_mae(
prediction, gt)
for pidx, pdata in enumerate(zip(precision, recall)):
p, r = pdata
precision_record[pidx].update(p)
recall_record[pidx].update(r)
mae_record.update(mae)
time_record.update(end - start)
if args['save_results']:
Image.fromarray(prediction).save(
os.path.join(
ckpt_path, exp_name,
'(%s) %s_%s' % (exp_name, name, args['snapshot']),
img_name + '.png'))
max_fmeasure, mean_fmeasure = cal_fmeasure_both(
[precord.avg for precord in precision_record],
[rrecord.avg for rrecord in recall_record])
results[name] = {
'max_fmeasure': max_fmeasure,
'mae': mae_record.avg,
'mean_fmeasure': mean_fmeasure
}
print 'test results:'
print results
with open('Result', 'a') as f:
if args['crf_refine']:
f.write('with CRF')
f.write('Runing time %.6f \n' % time_record.avg)
f.write('\n%s\n iter:%s\n' % (exp_name, args['snapshot']))
for name, value in results.iteritems():
f.write(
'%s: mean_fmeasure: %.10f, mae: %.10f, max_fmeasure: %.10f\n'
% (name, value['mean_fmeasure'], value['mae'],
value['max_fmeasure']))
