91 строка
3.4 KiB
Python
91 строка
3.4 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# Licensed under the MIT License.
|
|
|
|
import os
|
|
import sys
|
|
import torch
|
|
from torchvision.utils import save_image
|
|
from options.train_options import TrainOptions
|
|
import data
|
|
from util.iter_counter import IterationCounter
|
|
from util.util import print_current_errors
|
|
from util.util import mkdir
|
|
from trainers.pix2pix_trainer import Pix2PixTrainer
|
|
|
|
|
|
if __name__ == '__main__':
|
|
# parse options
|
|
opt = TrainOptions().parse()
|
|
# print options to help debugging
|
|
print(' '.join(sys.argv))
|
|
dataloader = data.create_dataloader(opt)
|
|
len_dataloader = len(dataloader)
|
|
# create tool for counting iterations
|
|
iter_counter = IterationCounter(opt, len(dataloader))
|
|
# create trainer for our model
|
|
trainer = Pix2PixTrainer(opt, resume_epoch=iter_counter.first_epoch)
|
|
save_root = os.path.join('checkpoints', opt.name, 'train')
|
|
mkdir(save_root)
|
|
|
|
for epoch in iter_counter.training_epochs():
|
|
opt.epoch = epoch
|
|
iter_counter.record_epoch_start(epoch)
|
|
for i, data_i in enumerate(dataloader, start=iter_counter.epoch_iter):
|
|
iter_counter.record_one_iteration()
|
|
# Training
|
|
# train generator
|
|
if i % opt.D_steps_per_G == 0:
|
|
trainer.run_generator_one_step(data_i)
|
|
# train discriminator
|
|
trainer.run_discriminator_one_step(data_i)
|
|
if iter_counter.needs_printing():
|
|
losses = trainer.get_latest_losses()
|
|
try:
|
|
print_current_errors(opt, epoch, iter_counter.epoch_iter,
|
|
iter_counter.epoch_iter_num, losses, iter_counter.time_per_iter)
|
|
except OSError as err:
|
|
print(err)
|
|
|
|
if iter_counter.needs_displaying():
|
|
imgs_num = data_i['label'].shape[0]
|
|
|
|
if opt.dataset_mode == 'deepfashionHD':
|
|
label = data_i['label'][:,:3,:,:]
|
|
|
|
show_size = opt.display_winsize
|
|
|
|
imgs = torch.cat((label.cpu(), data_i['ref'].cpu(), \
|
|
trainer.get_latest_generated().data.cpu(), \
|
|
data_i['image'].cpu()), 0)
|
|
|
|
try:
|
|
save_name = '%08d_%08d.png' % (epoch, iter_counter.total_steps_so_far)
|
|
save_name = os.path.join(save_root, save_name)
|
|
save_image(imgs, save_name, nrow=imgs_num, padding=0, normalize=True)
|
|
except OSError as err:
|
|
print(err)
|
|
|
|
if iter_counter.needs_saving():
|
|
print('saving the latest model (epoch %d, total_steps %d)' %
|
|
(epoch, iter_counter.total_steps_so_far))
|
|
try:
|
|
trainer.save('latest')
|
|
iter_counter.record_current_iter()
|
|
except OSError as err:
|
|
import pdb; pdb.set_trace()
|
|
print(err)
|
|
|
|
trainer.update_learning_rate(epoch)
|
|
iter_counter.record_epoch_end()
|
|
|
|
if epoch % opt.save_epoch_freq == 0 or epoch == iter_counter.total_epochs:
|
|
print('saving the model at the end of epoch %d, iters %d' %
|
|
(epoch, iter_counter.total_steps_so_far))
|
|
try:
|
|
trainer.save('latest')
|
|
trainer.save(epoch)
|
|
except OSError as err:
|
|
print(err)
|
|
|
|
print('Training was successfully finished.')
|