def get_initialized_network(opt) -> Tuple[nn.Module]:
"""
:return: content and pose encoder, decoder and scene discriminator, with `utils.init_weights` applied
"""
if opt.image_width == 64:
import models.resnet_64 as resnet_models
import models.dcgan_64 as dcgan_models
import models.dcgan_unet_64 as dcgan_unet_models
import models.vgg_unet_64 as vgg_unet_models
elif opt.image_width == 128:
import models.resnet_128 as resnet_models
import models.dcgan_128 as dcgan_models
import models.dcgan_unet_128 as dcgan_unet_models
import models.vgg_unet_128 as vgg_unet_models
import models.classifiers as classifiers
# load models
if opt.content_model == 'dcgan_unet':
netEC = dcgan_unet_models.content_encoder(opt.content_dim,
opt.channels)
netD = dcgan_unet_models.decoder(opt.content_dim, opt.pose_dim,
opt.channels)
elif opt.content_model == 'vgg_unet':
netEC = vgg_unet_models.content_encoder(opt.content_dim, opt.channels)
netD = vgg_unet_models.decoder(opt.content_dim, opt.pose_dim,
opt.channels)
elif opt.content_model == 'dcgan':
netEC = dcgan_models.content_encoder(opt.content_dim, opt.channels)
netD = dcgan_models.decoder(opt.content_dim, opt.pose_dim,
opt.channels)
else:
raise ValueError('Unknown content model: %s' % opt.content_model)
if opt.pose_model == 'dcgan':
netEP = dcgan_models.pose_encoder(opt.pose_dim,
opt.channels,
normalize=opt.normalize)
elif opt.pose_model == 'resnet':
netEP = resnet_models.pose_encoder(opt.pose_dim,
opt.channels,
normalize=opt.normalize)
else:
raise ValueError('Unknown pose model: %s' % opt.pose_model)
netC = classifiers.scene_discriminator(opt.pose_dim, opt.sd_nf)
netEC.apply(init_weights)
netEP.apply(init_weights)
netD.apply(init_weights)
netC.apply(init_weights)
return netEC, netEP, netD, netC
elif opt.image_width == 128:
import models.dcgan_128 as models
if opt.model == 'resnet':
if opt.image_width == 64:
raise ValueError('resnet_64 not implemented yet!')
elif opt.image_width == 128:
import models.resnet_128 as models
elif opt.model == 'unet':
if opt.image_width == 64:
import models.unet_64 as models
elif opt.image_width == 128:
raise ValueError('unet_128 not implemented yet!')
netC = models.scene_discriminator(opt.pose_dim)
netEC = models.content_encoder(opt.content_dim, opt.channels)
netEP = models.pose_encoder(opt.pose_dim, opt.channels)
netD = models.decoder(opt.content_dim, opt.pose_dim, opt.channels)
# ---------------- optimizers ----------------
if opt.optimizer == 'adam':
opt.optimizer = optim.Adam
elif opt.optimizer == 'rmsprop':
opt.optimizer = optim.RMSprop
elif opt.optimizer == 'sgd':
opt.optimizer = optim.SGD
else:
raise ValueError('Unknown optimizer: %s' % opt.optimizer)
optimizerC = opt.optimizer(netC.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
import models.vgg_unet_128 as vgg_unet_models
if opt.content_model == 'dcgan_unet':
netEC = dcgan_unet_models.content_encoder(opt.content_dim, opt.channels)
netD = dcgan_unet_models.decoder(opt.content_dim, opt.pose_dim, opt.channels)
elif opt.content_model == 'vgg_unet':
netEC = vgg_unet_models.content_encoder(opt.content_dim, opt.channels)
netD = vgg_unet_models.decoder(opt.content_dim, opt.pose_dim, opt.channels)
elif opt.content_model == 'dcgan':
netEC = dcgan_models.content_encoder(opt.content_dim, opt.channels)
netD = dcgan_models.decoder(opt.content_dim, opt.pose_dim, opt.channels)
else:
raise ValueError('Unknown content model: %s' % opt.content_model)
if opt.pose_model == 'dcgan':
netEP = dcgan_models.pose_encoder(opt.pose_dim, opt.channels, normalize=opt.normalize)
elif opt.pose_model == 'resnet':
netEP = resnet_models.pose_encoder(opt.pose_dim, opt.channels, normalize=opt.normalize)
else:
raise ValueError('Unknown pose model: %s' % opt.pose_model)
import models.classifiers as classifiers
netC = classifiers.scene_discriminator(opt.pose_dim, opt.sd_nf)
netEC.apply(utils.init_weights)
netEP.apply(utils.init_weights)
netD.apply(utils.init_weights)
netC.apply(utils.init_weights)
# ---------------- optimizers ----------------
import models.dcgan_128 as model
elif opt.model == 'vgg':
if opt.image_width == 64:
import models.vgg_64 as model
elif opt.image_width == 128:
import models.vgg_128 as model
else:
raise ValueError('Unknown model: %s' % opt.model)
# define
frame_predictor = lstm_models.lstm((opt.factor+1)*opt.z_dim, opt.g_dim, opt.rnn_size, opt.predictor_rnn_layers, int(opt.batch_size/len(opt.gpu_ids)))
posterior_pose = lstm_models.gaussian_lstm(opt.g_dim+opt.factor*opt.z_dim, opt.z_dim, opt.rnn_size, opt.posterior_rnn_layers, int(opt.batch_size/len(opt.gpu_ids)))
prior = lstm_models.gaussian_lstm(opt.g_dim+opt.factor*opt.z_dim, opt.z_dim, opt.rnn_size, opt.prior_rnn_layers, int(opt.batch_size/len(opt.gpu_ids)))
cont_encoder = model.cont_encoder(opt.z_dim*opt.factor, opt.channels*opt.n_past) #g_dim = 64 or 128
pose_encoder = model.pose_encoder(opt.g_dim, opt.channels)
decoder = model.decoder(opt.g_dim, opt.channels)
# init
frame_predictor = utils.init_net(frame_predictor, init_type='normal', init_gain=0.02, gpu_ids=opt.gpu_ids)
posterior_pose = utils.init_net(posterior_pose, init_type='normal', init_gain=0.02, gpu_ids=opt.gpu_ids)
prior = utils.init_net(prior, init_type='normal', init_gain=0.02, gpu_ids=opt.gpu_ids)
cont_encoder = utils.init_net(cont_encoder, init_type='normal', init_gain=0.02, gpu_ids=opt.gpu_ids)
pose_encoder = utils.init_net(pose_encoder, init_type='normal', init_gain=0.02, gpu_ids=opt.gpu_ids)
decoder = utils.init_net(decoder, init_type='normal', init_gain=0.02, gpu_ids=opt.gpu_ids)
# load
utils.load_network(frame_predictor, 'frame_predictor', 'last', opt.model_path,device)
utils.load_network(posterior_pose, 'posterior_pose', 'last', opt.model_path,device)
utils.load_network(prior, 'prior', 'last', opt.model_path,device)
netEC = dcgan_unet_models.content_encoder(opt.content_dim, opt.channels)
netD = dcgan_unet_models.decoder(opt.content_dim, opt.pose_dim, opt.channels)
netDis = dcgan_unet_models.Discriminator()
elif opt.content_model == 'vgg_unet':
netEC = vgg_unet_models.content_encoder(opt.content_dim, opt.channels)
netD = vgg_unet_models.decoder(opt.content_dim, opt.pose_dim, opt.channels)
elif opt.content_model == 'dcgan':
netEC = dcgan_models.content_encoder(opt.content_dim, opt.channels)
netD = dcgan_models.decoder(opt.content_dim, opt.pose_dim, opt.channels)
else:
raise ValueError('Unknown content model: %s' % opt.content_model)
if opt.pose_model == 'dcgan':
netEP = dcgan_models.pose_encoder(opt.pose_dim, opt.channels, normalize=opt.normalize)
elif opt.pose_model == 'resnet':
netEP = resnet_models.pose_encoder(opt.pose_dim, opt.channels, normalize=opt.normalize)
else:
raise ValueError('Unknown pose model: %s' % opt.pose_model)
import models.classifiers as classifiers
netC = classifiers.scene_discriminator(opt.pose_dim, opt.sd_nf)
netEC.apply(utils.init_weights)
netEP.apply(utils.init_weights)
netD.apply(utils.init_weights)
netC.apply(utils.init_weights)
netDis.apply(utils.init_weights)
# ---------------- optimizers ----------------