summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--train.py8
1 files changed, 4 insertions, 4 deletions
diff --git a/train.py b/train.py
index da497ce..ae129b9 100644
--- a/train.py
+++ b/train.py
@@ -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)' %