align_corners=True)
pred_post_init = generator.forward(images)
sal_loss = structure_loss(pred_post_init, gts)
sal_loss.backward()
generator_optimizer.step()
visualize_prediction_init(torch.sigmoid(pred_post_init))
visualize_gt(gts)
if rate == 1:
loss_record.update(sal_loss.data, opt.batchsize)
if i % 10 == 0 or i == total_step:
print(
'{} Epoch [{:03d}/{:03d}], Step [{:04d}/{:04d}], Gen Loss: {:.4f}'
.format(datetime.now(), epoch, opt.epoch, i, total_step,
loss_record.show()))
adjust_lr(generator_optimizer, opt.lr_gen, epoch, opt.decay_rate,
opt.decay_epoch)
save_path = 'models/Resnet/'
if not os.path.exists(save_path):
os.makedirs(save_path)
if epoch % opt.epoch == 0:
torch.save(generator.state_dict(),
save_path + 'Model' + '_%d' % epoch + '_gen.pth')
fixs = F.upsample(fixs, size=(trainsize, trainsize), mode='bilinear', align_corners=True)
fix_pred,cod_pred1,cod_pred2 = generator.forward(images)
fix_loss = mse_loss(torch.sigmoid(fix_pred),fixs)
cod_loss1 = structure_loss(cod_pred1, gts)
cod_loss2 = structure_loss(cod_pred2, gts)
final_loss = fix_loss + cod_loss1 + cod_loss2
final_loss.backward()
generator_optimizer.step()
visualize_cod1(torch.sigmoid(cod_pred1))
visualize_cod2(torch.sigmoid(cod_pred2))
visualize_fix(torch.sigmoid(fix_pred))
visualize_fix_gt(fixs)
visualize_gt(gts)
if rate == 1:
loss_record.update(final_loss.data, opt.batchsize)
if i % 10 == 0 or i == total_step:
print('{} Epoch [{:03d}/{:03d}], Step [{:04d}/{:04d}], Gen Loss: {:.4f}'.
format(datetime.now(), epoch, opt.epoch, i, total_step, loss_record.show()))
adjust_lr(generator_optimizer, opt.lr_gen, epoch, opt.decay_rate, opt.decay_epoch)
save_path = 'models/Resnet/'
if not os.path.exists(save_path):
os.makedirs(save_path)
if epoch % opt.epoch == 0:
torch.save(generator.state_dict(), save_path + 'Model' + '_%d' % epoch + '_gen.pth')