67 lines
2.6 KiB
Python
67 lines
2.6 KiB
Python
import os
|
|
import cv2
|
|
import numpy as np
|
|
import torchvision
|
|
from tensorboardX import SummaryWriter
|
|
|
|
|
|
class Saver():
|
|
def __init__(self, opts):
|
|
self.log_dir = opts.log_dir
|
|
self.image_dir = opts.img_dir
|
|
self.model_dir = opts.model_dir
|
|
|
|
self.log_freq = opts.log_freq
|
|
self.img_save_freq = opts.img_save_freq
|
|
self.model_save_freq = opts.model_save_freq
|
|
|
|
# create tensorboard writer
|
|
self.writer = SummaryWriter(logdir=self.log_dir)
|
|
|
|
# write losses and images to tensorboard
|
|
def write_log(self, total_it, model):
|
|
if not os.path.exists(self.log_dir):
|
|
os.makedirs(self.log_dir)
|
|
if (total_it + 1) % self.log_freq == 0:
|
|
# write loss
|
|
members = [attr for attr in dir(model) if
|
|
not callable(getattr(model, attr)) and not attr.startswith("__") and 'loss' in attr]
|
|
for m in members:
|
|
self.writer.add_scalar(m, getattr(model, m), total_it)
|
|
|
|
# save result images
|
|
def write_img(self, ep, model):
|
|
if not os.path.exists(self.image_dir):
|
|
os.makedirs(self.image_dir)
|
|
if (ep + 1) % self.img_save_freq == 0:
|
|
assembled_images = model.assemble_outputs()
|
|
img_filename = '%s/self_%05d.jpg' % (self.image_dir, ep)
|
|
torchvision.utils.save_image(assembled_images / 2 + 0.5, img_filename, nrow=1)
|
|
elif ep == -1:
|
|
assembled_images = model.assemble_outputs()
|
|
img_filename = '%s/self_last.jpg' % (self.image_dir)
|
|
torchvision.utils.save_image(assembled_images / 2 + 0.5, img_filename, nrow=1)
|
|
|
|
# save model
|
|
def write_model(self, ep, total_it, model):
|
|
if not os.path.exists(self.model_dir):
|
|
os.makedirs(self.model_dir)
|
|
if (ep + 1) % self.model_save_freq == 0:
|
|
print('--- save the model @ ep %d ---' % (ep))
|
|
model.save('%s/makeup_%05d.pth' % (self.model_dir, ep), ep, total_it)
|
|
elif ep == -1:
|
|
model.save('%s/makeup_last.pth' % self.model_dir, ep, total_it)
|
|
|
|
# save test pair images
|
|
def write_test_pair_img(self, iter, model, data):
|
|
root = os.path.join(self.image_dir, 'test_pair')
|
|
if not os.path.exists(root):
|
|
os.makedirs(root)
|
|
import time
|
|
start_time=time.time()
|
|
test_pair_img = model.test_pair(data)
|
|
print(time.time()-start_time)
|
|
img_filename = '%s/gen_%05d.jpg' % (root, iter)
|
|
torchvision.utils.save_image(test_pair_img / 2 + 0.5, img_filename, nrow=1)
|
|
|
|
|