def __init__(self, args):
args.deterministic = True
if args.dataset in ["mnist", "dsprites"]:
content_encoder = models.content_encoder(
args.g_dims, nc=args.num_channels).cuda()
position_encoder = models.pose_encoder(
args.z_dims,
nc=args.num_channels,
normalize=args.normalize_position).cuda()
else:
content_encoder = vgg_64.encoder(args.g_dims,
nc=args.num_channels).cuda()
position_encoder = resnet_64.pose_encoder(
args.z_dims, nc=args.num_channels).cuda()
if args.dataset == "mpi3d_real":
decoder = vgg_64.drnet_decoder(args.g_dims,
args.z_dims,
nc=args.num_channels).cuda()
else:
decoder = models.decoder(args.g_dims,
args.z_dims,
nc=args.num_channels,
skips=args.skips).cuda()
self.content_frames = 1
if args.content_lstm:
content_encoder = models.content_encoder_lstm(
args.g_dims, content_encoder, args.batch_size)
self.content_frames = args.input_frames
discriminator = models.scene_discriminator(args.z_dims).cuda()
nets = {
"content_encoder": content_encoder,
"position_encoder": position_encoder,
"decoder": decoder,
"discriminator": discriminator,
}
self.encoder_decoder_parameters = itertools.chain(*[
content_encoder.parameters(),
position_encoder.parameters(),
decoder.parameters(),
])
encoder_decoder_optim = torch.optim.Adam(
self.encoder_decoder_parameters,
lr=args.lr,
betas=(args.beta1, 0.999))
discriminator_optim = torch.optim.Adam(discriminator.parameters(),
lr=args.lr,
betas=(args.beta1, 0.999))
optims = {
"encoder_decoder_optim": encoder_decoder_optim,
"discriminator_optim": discriminator_optim,
}
super().__init__(nets, optims, args)
def __init__(self, z_dim, frame_encoding=128, nc=1, nout=1):
super().__init__()
self.frame_encoder = models.pose_encoder(frame_encoding, nc=nc)
self.z_embed_net = nn.Linear(z_dim, 128)
self.net_fc = nn.Sequential(nn.Linear(256,
64), nn.LeakyReLU(0.2, True),
nn.Linear(64, nout))
def __init__(self, args):
args.deterministic = True
encoder_checkpoint = torch.load(args.encoder_checkpoint)
if args.dataset in ["mnist", "dsprites"]:
Ec = models.content_encoder(args.g_dims,
nc=args.num_channels).cuda()
Ep = models.pose_encoder(args.z_dims, nc=args.num_channels).cuda()
else:
Ec = vgg_64.encoder(args.g_dims, nc=args.num_channels).cuda()
Ep = resnet_64.pose_encoder(args.z_dims,
nc=args.num_channels).cuda()
if args.dataset == "mpi3d_real":
D = vgg_64.drnet_decoder(args.g_dims,
args.z_dims,
nc=args.num_channels).cuda()
else:
D = models.decoder(args.g_dims,
args.z_dims,
nc=args.num_channels,
skips=args.skips).cuda()
Ep.load_state_dict(encoder_checkpoint["position_encoder"])
Ec.load_state_dict(encoder_checkpoint["content_encoder"])
D.load_state_dict(encoder_checkpoint["decoder"])
self.Ep = nn.DataParallel(Ep)
self.Ec = nn.DataParallel(Ec)
self.D = nn.DataParallel(D)
self.Ep.train()
self.Ec.train()
self.D.train()
lstm_model = lstm(args.g_dims + args.z_dims, args.z_dims,
args.rnn_size, args.rnn_layers,
args.batch_size).cuda()
nets = {"lstm": lstm_model}
lstm_optim = torch.optim.Adam(lstm_model.parameters(),
lr=args.lr,
betas=(args.beta1, 0.999))
optims = {"lstm_optim": lstm_optim}
super().__init__(nets, optims, args)
num_workers=5,
drop_last=True)
def get_data_batch():
while True:
for seq in data_loader:
seq[1].transpose_(2, 3).transpose_(1, 2)
yield seq
data_generator = get_data_batch()
if args.dataset == "dsprites":
Ec = models.content_encoder(args.g_dims, nc=args.num_channels).cuda()
Ep = models.pose_encoder(args.z_dims, nc=args.num_channels).cuda()
elif args.dataset == "mpi3d_real":
Ec = vgg_64.encoder(args.g_dims, nc=args.num_channels).cuda()
Ep = resnet_64.pose_encoder(args.z_dims, nc=args.num_channels).cuda()
checkpoint = torch.load(args.checkpoint)
Ec.load_state_dict(checkpoint["content_encoder"])
Ep.load_state_dict(checkpoint["position_encoder"])
Ec.eval()
Ep.eval()
latent_c = None
latent_p = None
factors_p = None
factors_c = None