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 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)
if opt.model_dir != '':
decoder = saved_model['decoder']
encoder = saved_model['encoder']
else:
encoder = model.encoder(opt.g_dim, opt.channels)
decoder = model.decoder(opt.g_dim, opt.channels)
encoder.apply(utils.init_weights)
decoder.apply(utils.init_weights)
frame_predictor_optimizer = opt.optimizer(frame_predictor.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
posterior_optimizer = opt.optimizer(posterior.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
prior_optimizer = opt.optimizer(prior.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
encoder_optimizer = opt.optimizer(encoder.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
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))
optimizerEC = opt.optimizer(netEC.parameters(),
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
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)
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)
