from custom_types import *
import options as options
from torch import nn
import models.factory as factory
from models.single_mesh_models import SingleMeshDiscriminator
from models.mesh_handler import MeshHandler, load_template_mesh
from dgts_base import DGTS
from process_data import mesh_utils
class Trainer(DGTS):
def __init__(self, opt: options.TrainOption, device: D):
super(Trainer, self).__init__(opt, device)
self.discriminator: SingleMeshDiscriminator = factory.model_lc(opt, SingleMeshDiscriminator, device=device)
self.optimizer_generator: Union[factory.OptimizerLC, N] = None
self.optimizer_discriminator: Union[factory.OptimizerLC, N] = None
self.template = MeshHandler(load_template_mesh(opt, opt.start_level)[1],
self.opt, self.opt.start_level).to(self.device)
self.real_mesh: Union[MeshHandler, N] = None
self.real_mesh_flipped: Union[MeshHandler, N] = None
self.mse = nn.MSELoss().to(device)
self.reconstruction_z = factory.NoiseMem(opt).load().to(device)
self.logger = factory.Logger(opt)
def get_z(self, reconstruction_mode: bool) -> List[Union[T, float]]:
if reconstruction_mode and self.level >= self.opt.reconstruction_start:
for i in range(len(self.reconstruction_z), self.level + 1):
self.reconstruction_z.append(0. if i else self.get_z_by_level(self.template, 0).detach())
return self.reconstruction_z
else:
return self.get_z_sequence(self.template, self.level)
def generate(self, reconstruction_mode: bool, inside_out: bool = False) -> MeshHandler:
z = self.get_z(reconstruction_mode)
template = self.template.copy()
if inside_out:
template = template.flip()
return self.generator(template, z, self.level)
def before_level(self) -> [Optimizer, Optimizer, MeshHandler]:
self.logger.start()
self.generator.dup(self.level)
self.discriminator.dup(self.level)
self.optimizer_generator = factory.OptimizerLC(self.opt, 'generator', self.generator.get_level(self.level))
self.optimizer_discriminator = factory.OptimizerLC(self.opt, 'discriminator', self.discriminator.get_level(self.level))
self.real_mesh = MeshHandler(mesh_utils.load_real_mesh(self.opt.mesh_name,self.opt.start_level + self.level),
self.opt, self.opt.start_level + self.level).to(self.device)
if self.opt.inside_out:
self.real_mesh_flipped = self.real_mesh.copy().flip()
def between_levels(self):
self.reconstruction_z.save()
self.generator.save()
self.discriminator.save()
self.opt.save()
self.logger.stop()
self.level += 1
def gradient_penalty(self, fake_mesh: MeshHandler) -> T:
fake_data = fake_mesh().data
real_data = self.real_mesh().data
alpha = torch.rand(1, device=self.device)
interpolates = alpha * real_data + ((1 - alpha) * fake_data)
interpolates = torch.autograd.Variable(interpolates, requires_grad=True)
disc_interpolates = self.discriminator.penalty_forward(self.real_mesh, self.level, interpolates)
gradients = torch.autograd.grad(outputs=disc_interpolates, inputs=interpolates,
grad_outputs=torch.ones(disc_interpolates.size(), device=self.device),
create_graph=True, retain_graph=True, only_inputs=True)[0]
gradient_penalty = (gradients.norm(2, dim=1) - 1) ** 2
return gradient_penalty.mean()
def train_level(self):
def penalty_iter():
self.optimizer_discriminator.zero_grad()
fake_mesh = self.generate(False, False).detach()
gradient_penalty = self.gradient_penalty(fake_mesh)
(self.opt.penalty_weight * gradient_penalty).backward(retain_graph=True)
self.optimizer_discriminator.step()
def discriminator_iter(inside_out: bool = False):
if inside_out and not self.opt.inside_out:
return
self.optimizer_discriminator.zero_grad()
fake_mesh = self.generate(False, inside_out).detach()
real_mesh = self.real_mesh_flipped if inside_out else self.real_mesh
out_real = self.discriminator(real_mesh.copy(), self.level)
error_real = out_real.mean()
out_fake = self.discriminator(fake_mesh, self.level)
error_fake = out_fake.mean()
gradient_penalty = self.gradient_penalty(fake_mesh)
(self.opt.penalty_weight * gradient_penalty).backward(retain_graph=True)
(error_fake - error_real).backward()
self.optimizer_discriminator.step()
self.logger.stash_iter('d_fake', error_fake, 'd_real', error_real)
def generator_iter(inside_out: bool = False):
if inside_out and not self.opt.inside_out:
return
nonlocal meshes
self.optimizer_generator.zero_grad()
fake_mesh = self.generate(False, inside_out)
rec_mesh = self.generate(True, inside_out)
out_fake = self.discriminator(fake_mesh, self.level)
error_fake = out_fake.mean()
error_rec = self.mse(rec_mesh.vs, self.real_mesh.vs)
rec_weight = self.opt.reconstruction_weight
fake_loss = - error_fake
(rec_weight * error_rec + fake_loss).backward()
self.optimizer_generator.step()
self.logger.stash_iter('g_fake', error_fake, 'g_rec', error_rec)
meshes = rec_mesh, fake_mesh
def train_iter():
for _ in range(self.opt.discriminator_iters):
discriminator_iter()
for _ in range(self.opt.generator_iters):
generator_iter()
self.logger.reset_iter()
def decay(self):
self.optimizer_generator.decay()
self.optimizer_discriminator.decay()
def mesh_paths(level: int, iteration: int) -> List[List[str]]:
return [[f'{self.opt.cp_folder}/{export_type}/{tag}_{level:02d}_{iteration + 1:04d}' for
export_type in ['generated', 'plots']] for tag in ['rec', 'fake']]
def plot(meshes: Tuple[MeshHandler, MeshHandler], paths: List[List[str]]):
for mesh, path in zip(meshes, paths):
mesh.export(path[0])
mesh.plot(path[1])
meshes: Union[Tuple[MeshHandler, ...], N] = None
for iteration in range(self.opt.level_iters[self.level + self.opt.start_level]):
train_iter()
factory.do_when_its_time(self.opt.lr_decay_every, decay, iteration, self)
factory.do_when_its_time(self.opt.export_meshes_every, plot, iteration, meshes, mesh_paths(self.level, iteration))
def train(self):
for _ in range(self.opt.num_levels):
self.before_level()
self.train_level()
self.between_levels()
if __name__ == '__main__':
opt_ = options.TrainOption()
opt_.parse_cmdline()
opt_ = opt_.load()
device_ = CUDA(0)
trainer = Trainer(opt_, device_)
trainer.train()
