def create_approximate_posterior():
if args.approximate_posterior_type == 'diagonal-normal':
context_encoder = nn_.ConvEncoder(
context_features=args.context_features,
channels_multiplier=16,
dropout_probability=args.dropout_probability_encoder_decoder)
approximate_posterior = distributions_.ConditionalDiagonalNormal(
shape=[args.latent_features], context_encoder=context_encoder)
else:
context_encoder = nn.Linear(args.context_features,
2 * args.latent_features)
distribution = distributions_.ConditionalDiagonalNormal(
shape=[args.latent_features], context_encoder=context_encoder)
transform = transforms.CompositeTransform([
transforms.CompositeTransform([
create_linear_transform(),
create_base_transform(
i, context_features=args.context_features)
]) for i in range(args.num_flow_steps)
])
transform = transforms.CompositeTransform(
[transform, create_linear_transform()])
approximate_posterior = flows.Flow(
transforms.InverseTransform(transform), distribution)
return approximate_posterior
def create_flow(c, h, w, flow_checkpoint, _log):
distribution = distributions.StandardNormal((c * h * w, ))
transform = create_transform(c, h, w)
flow = flows.Flow(transform, distribution)
_log.info('There are {} trainable parameters in this model.'.format(
utils.get_num_parameters(flow)))
if flow_checkpoint is not None:
flow.load_state_dict(torch.load(flow_checkpoint))
_log.info('Flow state loaded from {}'.format(flow_checkpoint))
return flow
def create_prior():
if args.prior_type == 'standard-normal':
prior = distributions_.StandardNormal((args.latent_features, ))
else:
distribution = distributions_.StandardNormal(
(args.latent_features, ))
transform = transforms.CompositeTransform([
transforms.CompositeTransform(
[create_linear_transform(),
create_base_transform(i)])
for i in range(args.num_flow_steps)
])
transform = transforms.CompositeTransform(
[transform, create_linear_transform()])
prior = flows.Flow(transform, distribution)
return prior
raise ValueError
def create_transform():
transform = transforms.CompositeTransform([
transforms.CompositeTransform(
[create_linear_transform(),
create_base_transform(i)]) for i in range(args.num_flow_steps)
] + [create_linear_transform()])
return transform
# create model
distribution = distributions.StandardNormal((features, ))
transform = create_transform()
flow = flows.Flow(transform, distribution).to(device)
n_params = utils.get_num_parameters(flow)
print('There are {} trainable parameters in this model.'.format(n_params))
# create optimizer
optimizer = optim.Adam(flow.parameters(), lr=args.learning_rate)
if args.anneal_learning_rate:
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer,
args.num_training_steps,
0)
else:
scheduler = None
# create summary writer and write to log directory
timestamp = cutils.get_timestamp()
apply_unconditional_transform=args.apply_unconditional_transform)
return base_transform
transform = transforms.CompositeTransform([
transforms.CompositeTransform(
[create_linear_transform(),
create_base_transform(i)]) for i in range(args.num_flow_steps)
] + [create_linear_transform()])
device = torch.device('cpu')
torch.set_default_tensor_type('torch.DoubleTensor')
base_dist = distributions.TweakedUniform(torch.tensor([-np.pi] * feature_dim),
torch.tensor([np.pi] * feature_dim))
flow = flows.Flow(transform, base_dist).to(device)
num_steps = 3000
data = torch.load("./output/flow_step_{}.pt".format(num_steps),
map_location=torch.device('cpu'))
flow.load_state_dict(data['state_dict'])
flow.eval()
## compute pq_logp and pq_logq
pq_logp = vmf.log_prob(xyz).numpy() + np.log(np.abs(np.cos(feature[:, 0])))
with torch.no_grad():
pq_logq = flow.log_prob(torch.from_numpy(feature), None).numpy()
## compute qp_logq and qp_logp
N = feature.shape[0]
z = base_dist.sample(N, None)