def build_net(self): """Setup generator, optimizer, loss func and transfer to device """ self.render_layer = renderer.RenderLayerPointLightEnvTorch() # Build net self.encoder = nn.DataParallel(network.encoderInitial(7), device_ids=self.opts.gpu_id).cuda() self.decoder_brdf = nn.DataParallel( network.decoderBRDF(), device_ids=self.opts.gpu_id).cuda() self.decoder_render = nn.DataParallel( network.decoderRender(litc=30), device_ids=self.opts.gpu_id).cuda() self.env_predictor = nn.DataParallel( network.envmapInitial(), device_ids=self.opts.gpu_id).cuda() # Optimizer self.optimizerE = torch.optim.Adam(self.encoder.parameters(), lr=1e-4, betas=(0.5, 0.999)) self.optimizerBRDF = torch.optim.Adam(self.decoder_brdf.parameters(), lr=2e-4, betas=(0.5, 0.999)) self.optimizerDRen = torch.optim.Adam(self.decoder_render.parameters(), lr=2e-4, betas=(0.5, 0.999)) self.optimizerEnv = torch.optim.Adam(self.env_predictor.parameters(), lr=2e-4, betas=(0.5, 0.999)) self.error_list_albedo = [] self.error_list_normal = [] self.error_list_depth = [] self.error_list_rough = [] self.error_list_env = [] self.error_list_relit = [] self.error_list_total = [] if self.opts.reuse: print('--> loading saved models and loss npys') [self.update_lr() for i in range(int(self.opts.start_epoch / 2))] self.load_saved_loss(self.opts.start_epoch) self.load_saved_checkpoint(self.opts.start_epoch) else: # loss for saving self.error_save_albedo = [] self.error_save_normal = [] self.error_save_depth = [] self.error_save_rough = [] self.error_save_env = [] self.error_save_relit = [] self.error_save_total = [] print('--> start a new model')
def build_net(self): """Setup generator, optimizer, loss func and transfer to device """ self.render_layer = renderer.RenderLayerPointLightEnvTorch() # Build net self.encoder = nn.DataParallel(network.encoderInitial(7), device_ids=self.opts.gpu_id).cuda() self.decoder_brdf = nn.DataParallel( network.decoderBRDF(), device_ids=self.opts.gpu_id).cuda() self.decoder_render = nn.DataParallel( network.decoderRender(litc=30), device_ids=self.opts.gpu_id).cuda() self.env_predictor = nn.DataParallel( network.envmapInitial(), device_ids=self.opts.gpu_id).cuda() self.encoderRef = nn.DataParallel(network.RefineEncoder(), device_ids=self.opts.gpu_id).cuda() self.decoderRef_brdf = nn.DataParallel( network.RefineDecoderBRDF(), device_ids=self.opts.gpu_id).cuda() self.decoderRef_render = nn.DataParallel( network.RefineDecoderRender(litc=30), device_ids=self.opts.gpu_id).cuda() self.env_caspredictor = nn.DataParallel( network.RefineDecoderEnv(), device_ids=self.opts.gpu_id).cuda() print('--> loading saved model') path = '%s/%s/state_dict_13/models' % (self.opts.outf, self.name) self.encoder.load_state_dict( torch.load('%s/encoder.pth' % path, map_location=lambda storage, loc: storage)) self.decoder_brdf.load_state_dict( torch.load('%s/decoder_brdf.pth' % path, map_location=lambda storage, loc: storage)) self.decoder_render.load_state_dict( torch.load('%s/decoder_render.pth' % path, map_location=lambda storage, loc: storage)) self.env_predictor.load_state_dict( torch.load('%s/env_predictor.pth' % path, map_location=lambda storage, loc: storage)) path = '%s/%s/cas1/state_dict_9/models' % (self.opts.outf, self.name) self.encoderRef.load_state_dict( torch.load('%s/encoderRef.pth' % path, map_location=lambda storage, loc: storage)) self.decoderRef_brdf.load_state_dict( torch.load('%s/decoderRef_brdf.pth' % path, map_location=lambda storage, loc: storage)) self.decoderRef_render.load_state_dict( torch.load('%s/decoderRef_render.pth' % path, map_location=lambda storage, loc: storage)) self.env_caspredictor.load_state_dict( torch.load('%s/env_caspredictor.pth' % path, map_location=lambda storage, loc: storage)) def _fix(model): model.eval() for param in model.parameters(): param.requires_grad = False _fix(self.encoder) _fix(self.decoder_brdf) _fix(self.decoder_render) _fix(self.env_predictor) _fix(self.encoderRef) _fix(self.decoderRef_brdf) _fix(self.decoderRef_render) _fix(self.env_caspredictor) # Trainable self.encoderRef2 = nn.DataParallel(network.RefineEncoder(), device_ids=self.opts.gpu_id).cuda() self.decoderRef2_brdf = nn.DataParallel( network.RefineDecoderBRDF(), device_ids=self.opts.gpu_id).cuda() self.decoderRef2_render = nn.DataParallel( network.RefineDecoderRender(litc=30), device_ids=self.opts.gpu_id).cuda() self.env_cas2predictor = nn.DataParallel( network.RefineDecoderEnv(), device_ids=self.opts.gpu_id).cuda() assert self.train is True # Optimizer self.w_brdf_A = 1 self.w_brdf_N = 1 self.w_brdf_R = 0.5 self.w_brdf_D = 0.5 self.w_relit = 1 # Optimizer, actually only a group of optimizer self.optimizerE = torch.optim.Adam(self.encoderRef2.parameters(), lr=1e-4, betas=(0.5, 0.999)) self.optimizerBRDF = torch.optim.Adam( self.decoderRef2_brdf.parameters(), lr=2e-4, betas=(0.5, 0.999)) self.optimizerDRen = torch.optim.Adam( self.decoderRef2_render.parameters(), lr=2e-4, betas=(0.5, 0.999)) self.optimizerEnv = torch.optim.Adam( self.env_cas2predictor.parameters(), lr=2e-4, betas=(0.5, 0.999)) self.error_list_albedo = [] self.error_list_normal = [] self.error_list_depth = [] self.error_list_rough = [] self.error_list_env = [] self.error_list_relit = [] self.error_list_total = [] if self.opts.reuse: print('--> loading saved models and loss npys') [self.update_lr() for i in range(int(self.opts.start_epoch / 2))] self.load_saved_loss(self.opts.start_epoch) self.load_saved_checkpoint(self.opts.start_epoch) else: # loss for saving self.error_save_albedo = [] self.error_save_normal = [] self.error_save_depth = [] self.error_save_rough = [] self.error_save_env = [] self.error_save_relit = [] self.error_save_total = [] print('--> start a new model')
from PIL import Image # Options opts = test_options.TestOptions().parse() opts.name = 'pt' opts.outf = './data/output' opts.workers = 2 opts.batch_size = 1 opts.gpu_id = [0] print('--> pytorch can use %d GPUs' % (torch.cuda.device_count())) print('--> pytorch is using %d GPUs' % (len(opts.gpu_id))) print('--> GPU IDs:', opts.gpu_id) # Model encoder = nn.DataParallel(network.encoderInitial(4), device_ids=opts.gpu_id).cuda() decoder_brdf = nn.DataParallel(network.decoderBRDF(), device_ids=opts.gpu_id).cuda() decoder_render = nn.DataParallel(network.decoderRender(), device_ids=opts.gpu_id).cuda() render_layer = render.RenderLayerPointLightTorch() encoderRef = nn.DataParallel(network.RefineEncoder(), device_ids=opts.gpu_id).cuda() decoderRef_brdf = nn.DataParallel(network.RefineDecoderBRDF(), device_ids=opts.gpu_id).cuda() decoderRef_render = nn.DataParallel(network.RefineDecoderRender(), device_ids=opts.gpu_id).cuda()