def __init__(self, opt): assert opt.isTrain opt = copy.deepcopy(opt) if len(opt.gpu_ids) > 0: opt.gpu_ids = opt.gpu_ids[:1] self.gpu_ids = opt.gpu_ids super(SPADEModelModules, self).__init__() self.opt = opt self.model_names = ['G_student', 'G_teacher', 'D'] teacher_opt = self.create_option('teacher') self.netG_teacher = networks.define_G(opt.teacher_netG, gpu_ids=self.gpu_ids, opt=teacher_opt) student_opt = self.create_option('student') self.netG_student = networks.define_G(opt.student_netG, init_type=opt.init_type, init_gain=opt.init_gain, gpu_ids=self.gpu_ids, opt=student_opt) if hasattr(opt, 'distiller'): pretrained_opt = self.create_option('pretrained') self.netG_pretrained = networks.define_G(opt.pretrained_netG, gpu_ids=self.gpu_ids, opt=pretrained_opt) self.netD = networks.define_D(opt.netD, init_type=opt.init_type, init_gain=opt.init_gain, gpu_ids=self.gpu_ids, opt=opt) self.mapping_layers = ['head_0', 'G_middle_1', 'up_1'] self.netAs = nn.ModuleList() for i, mapping_layer in enumerate(self.mapping_layers): if mapping_layer != 'up_1': fs, ft = opt.student_ngf * 16, opt.teacher_ngf * 16 else: fs, ft = opt.student_ngf * 4, opt.teacher_ngf * 4 if hasattr(opt, 'distiller'): netA = nn.Conv2d(in_channels=fs, out_channels=ft, kernel_size=1) else: netA = SuperConv2d(in_channels=fs, out_channels=ft, kernel_size=1) networks.init_net(netA, opt.init_type, opt.init_gain, self.gpu_ids) self.netAs.append(netA) self.criterionGAN = GANLoss(opt.gan_mode) self.criterionFeat = nn.L1Loss() self.criterionVGG = VGGLoss() self.optimizers = [] self.netG_teacher.eval() self.config = None
def define_G(input_nc, output_nc, ngf, netG, norm='batch', use_dropout=False, init_type='normal', init_gain=0.02, gpu_ids=[]): """Create a 3D generator Parameters: input_nc (int) -- the number of channels in input images output_nc (int) -- the number of channels in output images ngf (int) -- the number of filters in the last conv layer netG (str) -- the architecture's name: resnet_9blocks | resnet_6blocks | unet_256 | unet_128 norm (str) -- the name of normalization layers used in the network: batch | instance | none use_dropout (bool) -- if use dropout layers. init_type (str) -- the name of our initialization method. init_gain (float) -- scaling factor for normal, xavier and orthogonal. gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2 Returns a generator Our current implementation provides two types of generators: U-Net: [unet_128] (for 128x128 input images) and [unet_256] (for 256x256 input images) The original U-Net paper: https://arxiv.org/abs/1505.04597 Resnet-based generator: [resnet_6blocks] (with 6 Resnet blocks) and [resnet_9blocks] (with 9 Resnet blocks) Resnet-based generator consists of several Resnet blocks between a few downsampling/upsampling operations. We adapt Torch code from Justin Johnson's neural style transfer project (https://github.com/jcjohnson/fast-neural-style). The generator has been initialized by <init_net>. It uses RELU for non-linearity. """ net = None norm_layer = get_norm_layer(norm_type=norm) if netG == 'resnet_9blocks': raise NotImplementedError # net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=9) elif netG == 'resnet_6blocks': raise NotImplementedError # net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=6) elif netG == 'unet_128': raise NotImplementedError # net = UnetGenerator(input_nc, output_nc, 7, ngf, norm_layer=norm_layer, use_dropout=use_dropout) elif netG == 'unet_256': net = UnetGenerator3d(input_nc, output_nc, 8, ngf, norm_layer=norm_layer, use_dropout=use_dropout) else: raise NotImplementedError('Generator model name [%s] is not recognized' % netG) return init_net(net, init_type, init_gain, gpu_ids)
def define_D(input_nc, ndf, netD, n_layers_D=3, norm='batch', init_type='normal', init_gain=0.02, gpu_ids=[]): """Create a discriminator Parameters: input_nc (int) -- the number of channels in input images ndf (int) -- the number of filters in the first conv layer netD (str) -- the architecture's name: basic | n_layers | pixel n_layers_D (int) -- the number of conv layers in the discriminator; effective when netD=='n_layers' norm (str) -- the type of normalization layers used in the network. init_type (str) -- the name of the initialization method. init_gain (float) -- scaling factor for normal, xavier and orthogonal. gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2 Returns a discriminator Our current implementation provides three types of discriminators: [basic]: 'PatchGAN' classifier described in the original pix2pix paper. It can classify whether 70×70 overlapping patches are real or fake. Such a patch-level discriminator architecture has fewer parameters than a full-image discriminator and can work on arbitrarily-sized images in a fully convolutional fashion. [n_layers]: With this mode, you cna specify the number of conv layers in the discriminator with the parameter <n_layers_D> (default=3 as used in [basic] (PatchGAN).) [pixel]: 1x1 PixelGAN discriminator can classify whether a pixel is real or not. It encourages greater color diversity but has no effect on spatial statistics. The discriminator has been initialized by <init_net>. It uses Leakly RELU for non-linearity. """ net = None norm_layer = get_norm_layer(norm_type=norm) if netD == 'basic': # default PatchGAN classifier net = NLayerDiscriminator3d(input_nc, ndf, n_layers=3, norm_layer=norm_layer) elif netD == 'n_layers': # more options net = NLayerDiscriminator3d(input_nc, ndf, n_layers_D, norm_layer=norm_layer) elif netD == 'pixel': # classify if each pixel is real or fake raise NotImplementedError # net = PixelDiscriminator(input_nc, ndf, norm_layer=norm_layer) else: raise NotImplementedError('Discriminator model name [%s] is not recognized' % net) return init_net(net, init_type, init_gain, gpu_ids)
def __init__(self, opt): assert opt.isTrain valid_netGs = [ 'munit', 'super_munit', 'super_mobile_munit', 'super_mobile_munit2', 'super_mobile_munit3' ] assert opt.teacher_netG in valid_netGs and opt.student_netG in valid_netGs super(BaseMunitDistiller, self).__init__(opt) self.loss_names = [ 'G_gan', 'G_rec_x', 'G_rec_c', 'G_rec_s', 'D_fake', 'D_real' ] if not opt.student_no_style_encoder: self.loss_names.append('G_rec_s') self.optimizers = [] self.image_paths = [] self.visual_names = ['real_A', 'Sfake_B', 'Tfake_B', 'real_B'] self.model_names = ['netG_student', 'netG_teacher', 'netD'] opt_teacher = self.create_option('teacher') self.netG_teacher = networks.define_G(opt.teacher_netG, init_type=opt.init_type, init_gain=opt.init_gain, gpu_ids=self.gpu_ids, opt=opt_teacher) opt_student = self.create_option('student') self.netG_student = networks.define_G(opt.student_netG, init_type=opt.init_type, init_gain=opt.init_gain, gpu_ids=self.gpu_ids, opt=opt_student) self.netD = networks.define_D(opt.netD, input_nc=opt.output_nc, init_type='normal', init_gain=opt.init_gain, gpu_ids=self.gpu_ids, opt=opt) if hasattr(opt, 'distiller'): self.netA = nn.Conv2d(in_channels=4 * opt.student_ngf, out_channels=4 * opt.teacher_ngf, kernel_size=1).to(self.device) else: self.netA = SuperConv2d(in_channels=4 * opt.student_ngf, out_channels=4 * opt.teacher_ngf, kernel_size=1).to(self.device) networks.init_net(self.netA) self.netG_teacher.eval() self.criterionGAN = GANLoss(opt.gan_mode).to(self.device) self.criterionRec = torch.nn.L1Loss() G_params = [] G_params.append(self.netG_student.parameters()) G_params.append(self.netA.parameters()) self.optimizer_G = torch.optim.Adam(itertools.chain(*G_params), lr=opt.lr, betas=(opt.beta1, 0.999), weight_decay=opt.weight_decay) self.optimizer_D = torch.optim.Adam(self.netD.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999), weight_decay=opt.weight_decay) self.optimizers.append(self.optimizer_G) self.optimizers.append(self.optimizer_D) self.eval_dataloader = create_eval_dataloader(self.opt, direction=opt.direction) self.inception_model, _, _ = create_metric_models(opt, device=self.device) self.npz = np.load(opt.real_stat_path) self.is_best = False
print('device: {}'.format(device)) # initial mesh mesh = Mesh(opts.initial_mesh, device=device, hold_history=True) # input point cloud input_xyz, input_normals = utils.read_pts(opts.input_pc) # normalize point cloud based on initial mesh input_xyz /= mesh.scale input_xyz += mesh.translations[None, :] input_xyz = torch.Tensor(input_xyz).type(options.dtype()).to(device)[None, :, :] input_normals = torch.Tensor(input_normals).type(options.dtype()).to(device)[None, :, :] part_mesh = PartMesh(mesh, num_parts=options.get_num_parts(len(mesh.faces)), bfs_depth=opts.overlap) print(f'number of parts {part_mesh.n_submeshes}') net, optimizer, rand_verts, scheduler = init_net(mesh, part_mesh, device, opts) for i in range(opts.iterations): num_samples = options.get_num_samples(i % opts.upsamp) if opts.global_step: optimizer.zero_grad() start_time = time.time() for part_i, est_verts in enumerate(net(rand_verts, part_mesh)): if not opts.global_step: optimizer.zero_grad() part_mesh.update_verts(est_verts[0], part_i) num_samples = options.get_num_samples(i % opts.upsamp) recon_xyz, recon_normals = sample_surface(part_mesh.main_mesh.faces, part_mesh.main_mesh.vs.unsqueeze(0), num_samples) # calc chamfer loss w/ normals recon_xyz, recon_normals = recon_xyz.type(options.dtype()), recon_normals.type(options.dtype()) xyz_chamfer_loss, normals_chamfer_loss = chamfer_distance(recon_xyz, input_xyz, x_normals=recon_normals, y_normals=input_normals,
def __init__(self, opt): assert opt.isTrain super(BaseResnetDistiller, self).__init__(opt) self.loss_names = ['G_gan', 'G_distill', 'G_recon', 'D_fake', 'D_real'] self.optimizers = [] self.image_paths = [] self.visual_names = ['real_A', 'Sfake_B', 'Tfake_B', 'real_B'] self.model_names = ['netG_student', 'netG_teacher', 'netD'] self.netG_teacher = networks.define_G(opt.input_nc, opt.output_nc, opt.teacher_ngf, opt.teacher_netG, opt.norm, opt.teacher_dropout_rate, opt.init_type, opt.init_gain, self.gpu_ids, opt=opt) self.netG_student = networks.define_G(opt.input_nc, opt.output_nc, opt.student_ngf, opt.student_netG, opt.norm, opt.student_dropout_rate, opt.init_type, opt.init_gain, self.gpu_ids, opt=opt) if getattr(opt, 'sort_channels', False) and opt.restore_student_G_path is not None: self.netG_student_tmp = networks.define_G(opt.input_nc, opt.output_nc, opt.student_ngf, opt.student_netG.replace('super_', ''), opt.norm, opt.student_dropout_rate, opt.init_type, opt.init_gain, self.gpu_ids, opt=opt) if hasattr(opt, 'distiller'): self.netG_pretrained = networks.define_G(opt.input_nc, opt.output_nc, opt.pretrained_ngf, opt.pretrained_netG, opt.norm, 0, opt.init_type, opt.init_gain, self.gpu_ids, opt=opt) if opt.dataset_mode == 'aligned': self.netD = networks.define_D(opt.input_nc + opt.output_nc, opt.ndf, opt.netD, opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids) elif opt.dataset_mode == 'unaligned': self.netD = networks.define_D(opt.output_nc, opt.ndf, opt.netD, opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids) else: raise NotImplementedError('Unknown dataset mode [%s]!!!' % opt.dataset_mode) self.netG_teacher.eval() self.criterionGAN = models.modules.loss.GANLoss(opt.gan_mode).to(self.device) if opt.recon_loss_type == 'l1': self.criterionRecon = torch.nn.L1Loss() elif opt.recon_loss_type == 'l2': self.criterionRecon = torch.nn.MSELoss() elif opt.recon_loss_type == 'smooth_l1': self.criterionRecon = torch.nn.SmoothL1Loss() elif opt.recon_loss_type == 'vgg': self.criterionRecon = models.modules.loss.VGGLoss(self.device) else: raise NotImplementedError('Unknown reconstruction loss type [%s]!' % opt.loss_type) if isinstance(self.netG_teacher, nn.DataParallel): self.mapping_layers = ['module.model.%d' % i for i in range(9, 21, 3)] else: self.mapping_layers = ['model.%d' % i for i in range(9, 21, 3)] self.netAs = [] self.Tacts, self.Sacts = {}, {} G_params = [self.netG_student.parameters()] for i, n in enumerate(self.mapping_layers): ft, fs = self.opt.teacher_ngf, self.opt.student_ngf if hasattr(opt, 'distiller'): netA = nn.Conv2d(in_channels=fs * 4, out_channels=ft * 4, kernel_size=1). \ to(self.device) else: netA = SuperConv2d(in_channels=fs * 4, out_channels=ft * 4, kernel_size=1). \ to(self.device) networks.init_net(netA) G_params.append(netA.parameters()) self.netAs.append(netA) self.loss_names.append('G_distill%d' % i) self.optimizer_G = torch.optim.Adam(itertools.chain(*G_params), lr=opt.lr, betas=(opt.beta1, 0.999)) self.optimizer_D = torch.optim.Adam(self.netD.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999)) self.optimizers.append(self.optimizer_G) self.optimizers.append(self.optimizer_D) self.eval_dataloader = create_eval_dataloader(self.opt, direction=opt.direction) block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[2048] self.inception_model = InceptionV3([block_idx]) self.inception_model.to(self.device) self.inception_model.eval() if 'cityscapes' in opt.dataroot: self.drn_model = DRNSeg('drn_d_105', 19, pretrained=False) util.load_network(self.drn_model, opt.drn_path, verbose=False) if len(opt.gpu_ids) > 0: self.drn_model.to(self.device) self.drn_model = nn.DataParallel(self.drn_model, opt.gpu_ids) self.drn_model.eval() self.npz = np.load(opt.real_stat_path) self.is_best = False
def __init__(self, opts, input_dim, output_dim, lambda_ms=None): super(BicycleGANAdaIN, self).__init__() self.isTrain = (opts.phase == 'train') self.gpu_ids = opts.gpu_ids self.device = torch.device('cuda:{}'.format( self.gpu_ids[0])) if self.gpu_ids else torch.device('cpu') self.nz = opts.nz # generator self.netG = networks.init_net(Gen(input_dim, output_dim, style_dim=opts.nz), init_type='xavier', init_gain=0.02, gpu_ids=self.gpu_ids) # discriminator if self.isTrain: self.netD = networks.define_D(output_dim, 64, netD='basic_256_multi', norm='instance', num_Ds=2, gpu_ids=self.gpu_ids) self.netD2 = networks.define_D(output_dim, 64, netD='basic_256_multi', norm='instance', num_Ds=2, gpu_ids=self.gpu_ids) # encoder self.netE = networks.define_E(output_dim, opts.nz, 64, netE=opts.bicycleE, norm='instance', vaeLike=True, gpu_ids=self.gpu_ids) # loss and optimizer and scheduler if self.isTrain: self.criterionGAN = networks.GANLoss(gan_mode=opts.gan_mode).to( self.device) self.criterionL1 = torch.nn.L1Loss() self.criterionZ = torch.nn.L1Loss() self.lambda_ms = 0. if lambda_ms is None else lambda_ms self.optimizer_G = torch.optim.Adam(self.netG.parameters(), lr=opts.lr, betas=(0.5, 0.999)) self.optimizer_E = torch.optim.Adam(self.netE.parameters(), lr=opts.lr, betas=(0.5, 0.999)) self.optimizer_D = torch.optim.Adam(self.netD.parameters(), lr=opts.lr, betas=(0.5, 0.999)) self.optimizer_D2 = torch.optim.Adam(self.netD2.parameters(), lr=opts.lr, betas=(0.5, 0.999)) self.optimizers = [ self.optimizer_G, self.optimizer_E, self.optimizer_D, self.optimizer_D2 ]
def __init__(self, opts): BaseModel.__init__(self, opts) lr = self.opt.lr self.model_names = ['G_A', 'G_B'] self.loss_names = [ 'd_total', 'g_total', 'g_rec_x_a', 'g_rec_x_b', 'g_rec_s_a', 'g_rec_s_b', 'g_rec_c_a', 'g_rec_c_b', 'g_adv_a', 'g_adv_b' ] self.visual_names = [] # Initiate the networks if self.isTrain: self.model_names = ['G_A', 'G_B', 'D_A', 'D_B'] self.netD_A = init_net(MsImageDis(self.opt.input_dim_a, self.opt.dis), init_type=self.opt.init, gpu_ids=self.gpu_ids) self.netD_B = init_net(MsImageDis(self.opt.input_dim_b, self.opt.dis), init_type=self.opt.init, gpu_ids=self.gpu_ids) self.netG_A = init_net(AdaINGen(self.opt.input_dim_a, self.opt.gen), init_type=self.opt.init, gpu_ids=self.gpu_ids) self.netG_B = init_net(AdaINGen(self.opt.input_dim_b, self.opt.gen), init_type=self.opt.init, gpu_ids=self.gpu_ids) self.instancenorm = nn.InstanceNorm2d(512, affine=False) self.style_dim = self.opt.gen['style_dim'] # fix the noise used in sampling display_size = self.opt.display_size self.s_a_fixed = torch.randn(display_size, self.style_dim, 1, 1).to(self.device) self.s_b_fixed = torch.randn(display_size, self.style_dim, 1, 1).to(self.device) if self.isTrain: # Setup the optimizers d_params = list(self.netD_A.parameters()) + list( self.netD_B.parameters()) g_params = list(self.netG_A.parameters()) + list( self.netG_B.parameters()) self.optimizer_D = torch.optim.Adam( [p for p in d_params if p.requires_grad], lr=lr, betas=(self.opt.beta1, self.opt.beta2), weight_decay=self.opt.weight_decay) self.optimizer_G = torch.optim.Adam( [p for p in g_params if p.requires_grad], lr=lr, betas=(self.opt.beta1, self.opt.beta2), weight_decay=self.opt.weight_decay) self.optimizer_names = ['optimizer_D', 'optimizer_G'] self.optimizers.append(self.optimizer_D) self.optimizers.append(self.optimizer_G) self.scheduler_D = get_scheduler(self.optimizer_D, self.opt) self.scheduler_G = get_scheduler(self.optimizer_G, self.opt) self.schedulers = [self.scheduler_D, self.scheduler_G] # Load VGG model if needed if (self.opt.vgg_w is not None) and self.opt.vgg_w > 0: self.vgg = load_vgg16(self.opt.vgg_model_path + '/models') self.vgg.eval() for param in self.vgg.parameters(): param.requires_grad = False
def main(): print( '############################### train.py ###############################' ) # Set random seed for reproducibility manual_seed = 999 # manualSeed = random.randint(1, 10000) # use if you want new results print("Random Seed: ", manual_seed) print() random.seed(manual_seed) torch.manual_seed(manual_seed) # Hyper parameters workers = 2 batch_size = 128 image_size = 128 nc = 3 in_ngc = 3 out_ngc = 3 in_ndc = in_ngc + out_ngc out_ndc = 1 ngf = 64 ndf = 32 sf = 100 # style factor for generator learning_rate = 0.0005 beta1 = 0.5 epochs = 100 gpu = True load_saved_model = False # print hyper parameters print(f'number of workers : {workers}') print(f'batch size : {batch_size}') print(f'image size : {image_size}') print(f'number of channels : {nc}') print(f'generator feature map size : {ngf}') print(f'discriminator feature map size : {ndf}') print(f'style factor : {sf}') print(f'learning rate : {learning_rate}') print(f'beta1 : {beta1}') print(f'epochs: {epochs}') print(f'GPU: {gpu}') print(f'load saved model: {load_saved_model}') print() # set up GPU device cuda = True if gpu and torch.cuda.is_available() else False # load CelebA dataset download_path = '/home/pbuddare/EEE_598/data/CelebA' # download_path = '/Users/prasanth/Academics/ASU/FALL_2019/EEE_598_CIU/data/Project/CelebA' data_loader_src = prepare_celeba_data(download_path, batch_size, image_size, workers) # load respective cartoon dataset download_path = '/home/pbuddare/EEE_598/data/Cartoon' # download_path = '/Users/prasanth/Academics/ASU/FALL_2019/EEE_598_CIU/data/Project/Cartoon' data_loader_tgt = prepare_cartoon_data(download_path, batch_size, image_size, workers) # show sample images show_images( next(iter(data_loader_src))[0], (8, 8), 16, 'Training images (Natural)', 'human_real') show_images( next(iter(data_loader_tgt))[0], (8, 8), 16, 'Training images (Cartoon)', 'cartoon_real') # create generator and discriminator networks generator = Generator(in_ngc, out_ngc, ngf) discriminator = Discriminator(in_ndc, ndf) if cuda: generator.cuda() discriminator.cuda() init_net(generator, gpu_ids=[0]) init_net(discriminator, gpu_ids=[0]) # loss function and optimizers criterion_GAN = nn.BCEWithLogitsLoss() criterion_L1 = nn.L1Loss() optimizer_g = optim.Adam(generator.parameters(), lr=learning_rate, betas=(beta1, 0.999)) optimizer_d = optim.Adam(discriminator.parameters(), lr=learning_rate, betas=(beta1, 0.999)) # Train GAN loss_G, loss_D = train_gan(data_loader_src, data_loader_tgt, generator, discriminator, criterion_GAN, criterion_L1, optimizer_d, optimizer_g, sf, batch_size, epochs, cuda) # save parameters current_time = str(datetime.datetime.now().time()).replace( ':', '').replace('.', '') + '.pth' g_path = './project_G_' + current_time d_path = './project_D_' + current_time torch.save(generator.state_dict(), g_path) torch.save(discriminator.state_dict(), d_path) # generate and display fake images test_imgs = next(iter(data_loader_src))[0] show_images(test_imgs, (8, 8), 16, 'Testing images (Natural)', 'human_real_test') test_imgs = test_imgs.cuda() if cuda else test_imgs fake_imgs = generator(test_imgs).detach() show_images(fake_imgs.cpu(), (8, 8), 16, 'Fake images (Cartoon)', 'cartoon_fake')
def start(self): opt = self.opt dataloader = self.dataloader model = self.model modules_on_one_gpu = getattr(model, 'modules_on_one_gpu', model) logger = self.logger if self.task == 'distill': shrink(model, opt) modules_on_one_gpu.netG_student = init_net( modules_on_one_gpu.netG_student, opt.init_type, opt.init_gain, []).to(model.device) if getattr(opt, 'prune_continue', False): model.load_networks(restore_pretrain=False) logger.print_info('All networks loaded.') model.print_networks() if 'spade' in self.opt.distiller: logger.print_info( f'netG student FLOPs: {mc.unwrap_model(modules_on_one_gpu.netG_student).n_macs}.' ) else: logger.print_info( f'netG student FLOPs: {mc.unwrap_model(modules_on_one_gpu.netG_student).n_macs}; down sampling: {mc.unwrap_model(modules_on_one_gpu.netG_student).down_sampling.n_macs}; features: {mc.unwrap_model(modules_on_one_gpu.netG_student).features.n_macs}; up sampling: {mc.unwrap_model(modules_on_one_gpu.netG_student).up_sampling.n_macs}.' ) if getattr(opt, 'prune_only', False): return start_epoch = opt.epoch_base end_epoch = opt.epoch_base + opt.nepochs + opt.nepochs_decay - 1 total_iter = opt.iter_base for epoch in range(start_epoch, end_epoch + 1): epoch_start_time = time.time() for i, data_i in enumerate(dataloader): iter_start_time = time.time() model.set_input(data_i) model.optimize_parameters(total_iter) if total_iter % opt.print_freq == 0: losses = model.get_current_losses() logger.print_current_errors(epoch, total_iter, losses, time.time() - iter_start_time) logger.plot(losses, total_iter) if total_iter % opt.save_latest_freq == 0 or total_iter == opt.iter_base: self.evaluate( epoch, total_iter, 'Saving the latest model (epoch %d, total_steps %d)' % (epoch, total_iter)) if getattr(model, 'is_best', False): model.save_networks('iter%d' % total_iter) model.save_networks('best') if getattr(model, 'is_best_A', False): model.save_networks('iter%d' % total_iter) model.save_networks('best_A') if getattr(model, 'is_best_B', False): model.save_networks('iter%d' % total_iter) model.save_networks('best_B') total_iter += 1 logger.print_info( 'End of epoch %d / %d \t Time Taken: %.2f sec' % (epoch, end_epoch, time.time() - epoch_start_time)) if epoch % opt.save_epoch_freq == 0 or epoch == end_epoch: self.evaluate( epoch, total_iter, 'Saving the model at the end of epoch %d, iters %d' % (epoch, total_iter)) model.save_networks(epoch) if getattr(model, 'is_best', False): model.save_networks('iter%d' % total_iter) model.save_networks('best') if getattr(model, 'is_best_A', False): model.save_networks('iter%d' % total_iter) model.save_networks('best_A') if getattr(model, 'is_best_B', False): model.save_networks('iter%d' % total_iter) model.save_networks('best_B') model.update_learning_rate(logger)
def run(self): # mesh = Mesh(opts.initial_mesh, device=device, hold_history=True) mesh = vtkMesh(self.initPoly, device=device, hold_history=True) # input point cloud input_xyz, input_normals = self.MakeInputData(self.targetPoly) # normalize point cloud based on initial mesh input_xyz /= mesh.scale input_xyz += mesh.translations[None, :] input_xyz = torch.Tensor(input_xyz).type( options.dtype()).to(device)[None, :, :] input_normals = torch.Tensor(input_normals).type( options.dtype()).to(device)[None, :, :] part_mesh = PartMesh(mesh, num_parts=options.get_num_parts(len(mesh.faces)), bfs_depth=opts.overlap) print(f'number of parts {part_mesh.n_submeshes}') net, optimizer, rand_verts, scheduler = init_net( mesh, part_mesh, device, opts) beamgap_loss = BeamGapLoss(device) if opts.beamgap_iterations > 0: print('beamgap on') beamgap_loss.update_pm( part_mesh, torch.cat([input_xyz, input_normals], dim=-1)) for i in range(opts.iterations): num_samples = options.get_num_samples(i % opts.upsamp) if opts.global_step: optimizer.zero_grad() start_time = time.time() for part_i, est_verts in enumerate(net(rand_verts, part_mesh)): if not opts.global_step: optimizer.zero_grad() part_mesh.update_verts(est_verts[0], part_i) num_samples = options.get_num_samples(i % opts.upsamp) recon_xyz, recon_normals = sample_surface( part_mesh.main_mesh.faces, part_mesh.main_mesh.vs.unsqueeze(0), num_samples) # calc chamfer loss w/ normals recon_xyz, recon_normals = recon_xyz.type( options.dtype()), recon_normals.type(options.dtype()) xyz_chamfer_loss, normals_chamfer_loss = chamfer_distance( recon_xyz, input_xyz, x_normals=recon_normals, y_normals=input_normals, unoriented=opts.unoriented) if (i < opts.beamgap_iterations) and (i % opts.beamgap_modulo == 0): loss = beamgap_loss(part_mesh, part_i) else: loss = (xyz_chamfer_loss + (opts.ang_wt * normals_chamfer_loss)) if opts.local_non_uniform > 0: loss += opts.local_non_uniform * local_nonuniform_penalty( part_mesh.main_mesh).float() loss.backward() if not opts.global_step: optimizer.step() scheduler.step() part_mesh.main_mesh.vs.detach_() if opts.global_step: optimizer.step() scheduler.step() end_time = time.time() if i % 1 == 0: print( f'{os.path.basename(opts.input_pc)}; iter: {i} out of: {opts.iterations}; loss: {loss.item():.4f};' f' sample count: {num_samples}; time: {end_time - start_time:.2f}' ) # mesh.export(os.path.join("./", f'recon_iter_{i}.obj')) self.backwarded.emit(mesh.vs) # if (i > 0 and (i + 1) % opts.upsamp == 0): # mesh = part_mesh.main_mesh # num_faces = int(np.clip(len(mesh.faces) * 1.5, len(mesh.faces), opts.max_faces)) # if num_faces > len(mesh.faces) or opts.manifold_always: # # up-sample mesh # mesh = utils.manifold_upsample(mesh, opts.save_path, Mesh, # num_faces=min(num_faces, opts.max_faces), # res=opts.manifold_res, simplify=True) # part_mesh = PartMesh(mesh, num_parts=options.get_num_parts(len(mesh.faces)), bfs_depth=opts.overlap) # print(f'upsampled to {len(mesh.faces)} faces; number of parts {part_mesh.n_submeshes}') # net, optimizer, rand_verts, scheduler = init_net(mesh, part_mesh, device, opts) # if i < opts.beamgap_iterations: # print('beamgap updated') # beamgap_loss.update_pm(part_mesh, input_xyz) with torch.no_grad(): mesh.export(os.path.join(opts.save_path, 'last_recon.obj'))
def __init__(self, opt): assert opt.isTrain valid_netGs = [ 'resnet_9blocks', 'mobile_resnet_9blocks', 'super_mobile_resnet_9blocks', 'sub_mobile_resnet_9blocks' ] assert opt.teacher_netG in valid_netGs and opt.student_netG in valid_netGs super(BaseResnetDistiller, self).__init__(opt) self.loss_names = ['G_gan', 'G_distill', 'G_recon', 'D_fake', 'D_real'] self.optimizers = [] self.image_paths = [] self.visual_names = ['real_A', 'Sfake_B', 'Tfake_B', 'real_B'] self.model_names = ['netG_student', 'netG_teacher', 'netD'] self.netG_teacher = networks.define_G( opt.teacher_netG, input_nc=opt.input_nc, output_nc=opt.output_nc, ngf=opt.teacher_ngf, norm=opt.norm, dropout_rate=opt.teacher_dropout_rate, gpu_ids=self.gpu_ids, opt=opt) self.netG_student = networks.define_G( opt.student_netG, input_nc=opt.input_nc, output_nc=opt.output_nc, ngf=opt.student_ngf, norm=opt.norm, dropout_rate=opt.student_dropout_rate, init_type=opt.init_type, init_gain=opt.init_gain, gpu_ids=self.gpu_ids, opt=opt) if hasattr(opt, 'distiller'): self.netG_pretrained = networks.define_G(opt.pretrained_netG, input_nc=opt.input_nc, output_nc=opt.output_nc, ngf=opt.pretrained_ngf, norm=opt.norm, gpu_ids=self.gpu_ids, opt=opt) if opt.dataset_mode == 'aligned': self.netD = networks.define_D(opt.netD, input_nc=opt.input_nc + opt.output_nc, ndf=opt.ndf, n_layers_D=opt.n_layers_D, norm=opt.norm, init_type=opt.init_type, init_gain=opt.init_gain, gpu_ids=self.gpu_ids, opt=opt) elif opt.dataset_mode == 'unaligned': self.netD = networks.define_D(opt.netD, input_nc=opt.output_nc, ndf=opt.ndf, n_layers_D=opt.n_layers_D, norm=opt.norm, init_type=opt.init_type, init_gain=opt.init_gain, gpu_ids=self.gpu_ids, opt=opt) else: raise NotImplementedError('Unknown dataset mode [%s]!!!' % opt.dataset_mode) self.netG_teacher.eval() self.criterionGAN = GANLoss(opt.gan_mode).to(self.device) if opt.recon_loss_type == 'l1': self.criterionRecon = torch.nn.L1Loss() elif opt.recon_loss_type == 'l2': self.criterionRecon = torch.nn.MSELoss() elif opt.recon_loss_type == 'smooth_l1': self.criterionRecon = torch.nn.SmoothL1Loss() elif opt.recon_loss_type == 'vgg': self.criterionRecon = models.modules.loss.VGGLoss(self.device) else: raise NotImplementedError( 'Unknown reconstruction loss type [%s]!' % opt.loss_type) if isinstance(self.netG_teacher, nn.DataParallel): self.mapping_layers = [ 'module.model.%d' % i for i in range(9, 21, 3) ] else: self.mapping_layers = ['model.%d' % i for i in range(9, 21, 3)] self.netAs = [] self.Tacts, self.Sacts = {}, {} G_params = [self.netG_student.parameters()] for i, n in enumerate(self.mapping_layers): ft, fs = self.opt.teacher_ngf, self.opt.student_ngf if hasattr(opt, 'distiller'): netA = nn.Conv2d(in_channels=fs * 4, out_channels=ft * 4, kernel_size=1). \ to(self.device) else: netA = SuperConv2d(in_channels=fs * 4, out_channels=ft * 4, kernel_size=1). \ to(self.device) networks.init_net(netA) G_params.append(netA.parameters()) self.netAs.append(netA) self.loss_names.append('G_distill%d' % i) self.optimizer_G = torch.optim.Adam(itertools.chain(*G_params), lr=opt.lr, betas=(opt.beta1, 0.999)) self.optimizer_D = torch.optim.Adam(self.netD.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999)) self.optimizers.append(self.optimizer_G) self.optimizers.append(self.optimizer_D) self.eval_dataloader = create_eval_dataloader(self.opt, direction=opt.direction) self.inception_model, self.drn_model, _ = create_metric_models( opt, device=self.device) self.npz = np.load(opt.real_stat_path) self.is_best = False