iter_start_time = time.time()
if total_steps % opt.print_freq == 0:
t_data = iter_start_time - iter_data_time
total_steps += opt.batch_size
epoch_iter += opt.batch_size
model.set_input(data)
model.optimize_parameters()
loss_mat.append(model.loss.cpu().data.numpy())
#CE_mat.append(model.CE_loss.cpu().data.numpy())
#prior_mat.append(model.prior_loss.cpu().data.numpy())
if total_steps % opt.print_freq == 0:
loss = model.loss
t = (time.time() - iter_start_time) / opt.batch_size
writer.print_current_losses(epoch, epoch_iter, loss, t, t_data)
writer.plot_loss(loss, epoch, epoch_iter, dataset_size)
if i % opt.save_latest_freq == 0:
print('saving the latest model (epoch %d, total_steps %d)' %
(epoch, total_steps))
model.save_network('latest')
iter_data_time = time.time()
writer.save_losses(loss_mat, CE_mat, prior_mat, epoch)
if epoch % opt.save_epoch_freq == 0:
print('saving the model at the end of epoch %d, iters %d' %
(epoch, total_steps))
model.save_network('latest')
model.save_network(epoch)
