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
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 ----------------
if opt.optimizer == 'adam':
opt.optimizer = optim.Adam
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 ----------------
if opt.optimizer == 'adam':
opt.optimizer = optim.Adam