71 lines
2.3 KiB
Python
71 lines
2.3 KiB
Python
import os
|
|
import torch
|
|
from .saver import Saver
|
|
from .options import Options
|
|
from .model import CSDMT
|
|
from .dataset import MakeupDataset
|
|
|
|
import warnings
|
|
warnings.filterwarnings("ignore")
|
|
|
|
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
|
|
def train():
|
|
# parse options
|
|
parser = Options()
|
|
opts = parser.parse()
|
|
|
|
# load dataset
|
|
print('\n--- load dataset ---')
|
|
dataset = MakeupDataset(opts)
|
|
train_loader = torch.utils.data.DataLoader(dataset, batch_size=opts.batch_size, shuffle=True,
|
|
num_workers=opts.nThreads)
|
|
print('dataset size:',len(dataset))
|
|
# load model
|
|
print('\n--- load model ---')
|
|
model = CSDMT(opts)
|
|
if opts.resume is None:
|
|
ep0 = -1
|
|
total_it = 0
|
|
else:
|
|
ep0, total_it = model.resume(opts.resume)
|
|
model.set_scheduler(opts, last_ep=ep0)
|
|
ep0 += 1
|
|
print('start the training at epoch %d' % (ep0))
|
|
# saver for output
|
|
saver = Saver(opts)
|
|
# train
|
|
print('\n--- train ---')
|
|
max_it = 1000000
|
|
for ep in range(ep0, opts.n_ep):
|
|
for it, data in enumerate(train_loader):
|
|
# load data to gpu
|
|
model.load_data(data)
|
|
# updata corr network
|
|
model.update_D()
|
|
model.update_G()
|
|
# print
|
|
if it%10==0:
|
|
print('total_it: %d (ep %d, it %d), lr %08f' % (total_it, ep, it, model.dis_opt.param_groups[0]['lr']))
|
|
print('G_loss: %.2f,loss_G_contrastive: %.2f,loss_G_GAN: %.2f,loss_G_selfL1: %.2f,loss_G_cycleL1: %.2f ' %
|
|
(model.G_loss,model.loss_G_contrastive,model.loss_G_GAN,model.loss_G_selfL1,model.loss_G_cycleL1 ))
|
|
total_it += 1
|
|
# save log
|
|
saver.write_log(total_it, model)
|
|
# save the last model
|
|
if total_it >= max_it:
|
|
saver.write_img(-1, model)
|
|
saver.write_model(-1, total_it,model)
|
|
break
|
|
# decay learning rate
|
|
if opts.n_ep_decay > -1:
|
|
model.update_lr()
|
|
# save result image
|
|
saver.write_img(ep, model)
|
|
# Save network weights
|
|
saver.write_model(ep, total_it, model)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
train()
|
|
print('The training is complete')
|