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()
