| import time |
| from options.train_options import TrainOptions |
| from data import CreateDataLoader |
| from models import create_model |
| from util.visualizer import Visualizer |
|
|
| if __name__ == '__main__': |
| start = time.time() |
| opt = TrainOptions().parse() |
| data_loader = CreateDataLoader(opt) |
| dataset = data_loader.load_data() |
| dataset_size = len(data_loader) |
| print('#training images = %d' % dataset_size) |
|
|
| model = create_model(opt) |
| model.setup(opt) |
| visualizer = Visualizer(opt) |
| total_steps = 0 |
| model.save_networks2(opt.which_epoch) |
|
|
| for epoch in range(opt.epoch_count, opt.niter + opt.niter_decay + 1): |
| epoch_start_time = time.time() |
| iter_data_time = time.time() |
| epoch_iter = 0 |
|
|
| for i, data in enumerate(dataset): |
| iter_start_time = time.time() |
| if total_steps % opt.print_freq == 0: |
| t_data = iter_start_time - iter_data_time |
| visualizer.reset() |
| total_steps += opt.batch_size |
| epoch_iter += opt.batch_size |
| model.set_input(data) |
| model.optimize_parameters() |
|
|
| if total_steps % opt.display_freq == 0: |
| save_result = total_steps % opt.update_html_freq == 0 |
| visualizer.display_current_results(model.get_current_visuals(), epoch, save_result) |
| |
|
|
| if total_steps % opt.print_freq == 0: |
| losses = model.get_current_losses() |
| t = (time.time() - iter_start_time) / opt.batch_size |
| visualizer.print_current_losses(epoch, epoch_iter, losses, t, t_data) |
| if opt.display_id > 0: |
| visualizer.plot_current_losses(epoch, float(epoch_iter) / dataset_size, opt, losses) |
| |
| if total_steps % opt.save_latest_freq == 0: |
| print('saving the latest model (epoch %d, total_steps %d)' % |
| (epoch, total_steps)) |
| |
| model.save_networks2('latest') |
|
|
| iter_data_time = time.time() |
| if epoch % opt.save_epoch_freq == 0: |
| print('saving the model at the end of epoch %d, iters %d' % |
| (epoch, total_steps)) |
| |
| |
| model.save_networks2('latest') |
| model.save_networks2(epoch) |
|
|
| print('End of epoch %d / %d \t Time Taken: %d sec' % |
| (epoch, opt.niter + opt.niter_decay, time.time() - epoch_start_time)) |
| model.update_learning_rate() |
|
|
| print('Total Time Taken: %d sec' % (time.time() - start)) |
|
|
