diff options
Diffstat (limited to 'train.py')
| -rw-r--r-- | train.py | 8 |
1 files changed, 4 insertions, 4 deletions
@@ -8,8 +8,8 @@ from util.visualizer import Visualizer data_loader = CreateDataLoader(opt) dataset = data_loader.load_data() -num_train = len(data_loader) -print('#training images = %d' % num_train) +dataset_size = len(data_loader) +print('#training images = %d' % dataset_size) model = create_model(opt) visualizer = Visualizer(opt) @@ -21,7 +21,7 @@ for epoch in range(1, opt.niter + opt.niter_decay + 1): for i, data in enumerate(dataset): iter_start_time = time.time() total_steps += opt.batchSize - epoch_iter = total_steps % num_train + epoch_iter = total_steps - dataset_size * (epoch - 1) model.set_input(data) model.optimize_parameters() @@ -32,7 +32,7 @@ for epoch in range(1, opt.niter + opt.niter_decay + 1): errors = model.get_current_errors() visualizer.print_current_errors(epoch, epoch_iter, errors, iter_start_time) if opt.display_id > 0: - visualizer.plot_current_errors(epoch, float(epoch_iter)/num_train, opt, errors) + visualizer.plot_current_errors(epoch, float(epoch_iter)/dataset_size, opt, errors) if total_steps % opt.save_latest_freq == 0: print('saving the latest model (epoch %d, total_steps %d)' % |