-
Notifications
You must be signed in to change notification settings - Fork 4
/
main.py
121 lines (103 loc) · 6.88 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
"""General-purpose training script for image-to-image translation.
This script works for various models (with option '--model': e.g., pix2pix, cyclegan, colorization) and
different datasets (with option '--dataset_mode': e.g., aligned, unaligned, single, colorization).
You need to specify the dataset ('--dataroot'), experiment name ('--name'), and model ('--model').
It first creates model, dataset, and visualizer given the option.
It then does standard network training. During the training, it also visualize/save the images, print/save the loss plot, and save models.
The script supports continue/resume training. Use '--continue_train' to resume your previous training.
Example:
Train a CycleGAN model:
python train.py --dataroot ./datasets/maps --name maps_cyclegan --model cycle_gan
Train a pix2pix model:
python train.py --dataroot ./datasets/facades --name facades_pix2pix --model pix2pix --direction BtoA
See options/base_options.py and options/train_options.py for more training options.
See training and test tips at: https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/docs/tips.md
See frequently asked questions at: https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/docs/qa.md
"""
import time
from options.options import Options
from data import create_dataset, get_test_loaders
from models import create_model
from models import misc
from util.visualizer import Visualizer
from util.util import format_time
from util import util
import tensorboardX as tensorboard
import os
import numpy as np
import torch
import copy
def test(model, opt, test_loader_a, test_loader_b):
# test
print('[*] testing start!')
dict_a, dict_b = misc.test_fid(test_loader_a, model.gen_a2b, test_loader_b, model.gen_b2a, model.run_dir, opt)
print('[Evaluation] A2B: ')
print(dict_a)
print('[Evaluation] B2A: ')
print(dict_b)
print('[*] testing finished!')
def training_loop(model, opt, dataset, test_loader_a, test_loader_b):
#-----------------------------------------------------------
# fix images for visualization
fix_a = torch.stack([test_loader_a.dataset[i]['A'] for i in range(opt.display_size)]).cuda() # fixed test data
fix_b = torch.stack([test_loader_b.dataset[i]['A'] for i in range(opt.display_size)]).cuda()
fix_train_a = torch.stack([dataset.dataset[i]['A'] for i in range(opt.batch_size)]).cuda() # fixed with different runs
fix_train_b = torch.stack([dataset.dataset[i]['B'] for i in range(opt.batch_size)]).cuda() # fixed with same run
#-----------------------------------------------------------
visualizer = Visualizer(opt) # create a visualizer that display/save images and plots
train_writer = tensorboard.SummaryWriter(os.path.join(model.run_dir, 'log')) # setup train writer
#-----------------------------------------------------------
cur_iters = opt.epoch_count * opt.iterations_per_epoch
used_time = opt.used_time
start_time = time.time() - used_time
total_iters = (opt.n_epochs + opt.n_epochs_decay) * opt.iterations_per_epoch
print('[*] training start!\n')
for epoch in range(opt.epoch_count+1, opt.n_epochs + opt.n_epochs_decay + 1): # outer loop for different epochs; we save the model by <epoch_count>, <epoch_count>+<save_latest_freq>
visualizer.reset() # reset the visualizer: make sure it saves the results to HTML at least once every epoch
for _ in range(opt.iterations_per_epoch//opt.batch_size): # inner loop within one epoch
data = dataset.next()
cur_iters += opt.batch_size
model.set_input(data) # unpack data from dataset and apply preprocessing
input_images, fake_images, betas = model.optimize_parameters() # calculate loss functions, get gradients, update network weights
if cur_iters % opt.print_freq == 0: # print training losses and save logging information to the disk
used_time = (time.time() - start_time)
util.write_loss(cur_iters, model, train_writer, prefix='training')
visualizer.print_current_losses(cur_iters, total_iters, used_time, model.get_current_losses())
if cur_iters % opt.display_freq == 0: # display images on visdom and save images to a HTML file
model.compute_visuals()
visualizer.display_current_results(model.get_current_visuals(), epoch, save_result=True)
# evaluation
test_images = model.test(fix_a, fix_b)
test_path = os.path.join(model.run_dir, 'img', 'fix-test_image-%03d.jpg' % epoch)
misc.save_image_grid(test_images, test_path, opt.display_size)
train_path = os.path.join(model.run_dir, 'img', 'fix-train_image-%03d.jpg' % epoch)
fix_images, fix_betas = model.get_betas(fix_train_a, fix_train_b)
misc.save_train_image_grid(fix_images, fix_betas, train_path)
train_path = os.path.join(model.run_dir, 'img', 'train_image-%03d.jpg' % epoch)
misc.save_train_image_grid(input_images, betas, train_path)
train_path = os.path.join(model.run_dir, 'img', 'train_image_fake-%03d.jpg' % epoch)
misc.save_train_image_grid(fake_images, betas, train_path)
model.update_learning_rate() # update learning rates in the beginning of every epoch.
if epoch % opt.save_epoch_freq == 0: # cache our model every <save_epoch_freq> epochs
model.save_networks('latest', used_time)
model.save_networks(epoch, used_time)
print('[*] training finished!')
if __name__ == '__main__':
opt = Options().parse() # get training options
#-----------------------------------------------------------
dataset = create_dataset(opt) # create a dataset given opt.dataset_mode and other options
opt.trainA_size = dataset.dataset.A_size # get the number of images in the dataset.
opt.trainB_size = dataset.dataset.B_size # get the number of images in the dataset.
test_loader_a, test_loader_b = get_test_loaders(opt) # get test loader by hard-coding options
opt.testA_size = len(test_loader_a.dataset) # get the number of images in the dataset.
opt.testB_size = len(test_loader_b.dataset) # get the number of images in the dataset.
#-----------------------------------------------------------
model = create_model(opt) # create a model given opt.model and other options
model.setup(opt) # regular setup: load and print networks; create schedulers; update opt.epoch_count
#-----------------------------------------------------------
if opt.phase == 'train' or opt.phase == 'resume':
training_loop(model, opt, dataset, test_loader_a, test_loader_b)
elif opt.phase == 'test':
test(model, opt, test_loader_a, test_loader_b)
else:
raise ValueError('opt.phase %s is not recognizable' % opt.phase)