def __init__(self, latent_size, data_size, flow_depth):
super().__init__()
hidden_size = latent_size * 2
self.inference_network = NeuralNetwork(input_size=data_size,
output_size=latent_size * 3,
hidden_size=hidden_size)
modules = []
for _ in range(flow_depth):
modules.append(flow.InverseAutoregressiveFlow(
num_input=latent_size, num_hidden=hidden_size,
num_context=latent_size))
modules.append(flow.Reverse(latent_size))
self.q_z_flow = flow.FlowSequential(*modules)
self.log_q_z_0 = NormalLogProb()
self.softplus = nn.Softplus()
def __init__(self, latent_size, flow_depth=2, logprob=False):
super().__init__()
if logprob:
self.encode_func = self.encode_logprob
else:
self.encode_func = self.encode
DIM = 64
self.main = nn.Sequential(
nn.Conv2d(1, DIM, 5, stride=2, padding=2),
nn.ReLU(True),
nn.Conv2d(DIM, 2 * DIM, 5, stride=2, padding=2),
nn.ReLU(True),
nn.Conv2d(2 * DIM, 4 * DIM, 5, stride=2, padding=2),
nn.ReLU(True),
)
if flow_depth > 0:
# IAF
hidden_size = latent_size * 2
flow_layers = [
flow.InverseAutoregressiveFlow(latent_size, hidden_size,
latent_size)
for _ in range(flow_depth)
]
flow_layers.append(flow.Reverse(latent_size))
self.q_z_flow = flow.FlowSequential(*flow_layers)
self.enc_chunk = 3
else:
self.q_z_flow = None
self.enc_chunk = 2
fc_out_size = latent_size * self.enc_chunk
conv_out_size = 4 * 4 * 4 * DIM
self.fc = nn.Sequential(
nn.Linear(conv_out_size, fc_out_size),
nn.LayerNorm(fc_out_size),
nn.LeakyReLU(0.2),
nn.Linear(fc_out_size, fc_out_size),
)
def __init__(self, cconv, latent_size, channels, flow_depth=2):
super().__init__()
self.cconv = cconv
if flow_depth > 0:
hidden_size = latent_size * 2
flow_layers = [
flow.InverseAutoregressiveFlow(latent_size, hidden_size,
latent_size)
for _ in range(flow_depth)
]
flow_layers.append(flow.Reverse(latent_size))
self.q_z_flow = flow.FlowSequential(*flow_layers)
self.enc_chunk = 3
else:
self.q_z_flow = None
self.enc_chunk = 2
self.grid_encoder = GridEncoder(channels, latent_size * self.enc_chunk)
def __init__(self, latent_size, flow_depth=2, logprob=False):
super().__init__()
if logprob:
self.encode_func = self.encode_logprob
else:
self.encode_func = self.encode
dim = 64
self.ls = nn.Sequential(nn.Conv2d(3, dim, 5, 2, 2), nn.LeakyReLU(0.2),
conv_ln_lrelu(dim, dim * 2),
conv_ln_lrelu(dim * 2, dim * 4),
conv_ln_lrelu(dim * 4, dim * 8),
nn.Conv2d(dim * 8, latent_size, 4))
if flow_depth > 0:
# IAF
hidden_size = latent_size * 2
flow_layers = [
flow.InverseAutoregressiveFlow(latent_size, hidden_size,
latent_size)
for _ in range(flow_depth)
]
flow_layers.append(flow.Reverse(latent_size))
self.q_z_flow = flow.FlowSequential(*flow_layers)
self.enc_chunk = 3
else:
self.q_z_flow = None
self.enc_chunk = 2
fc_out_size = latent_size * self.enc_chunk
self.fc = nn.Sequential(
nn.Linear(latent_size, fc_out_size),
nn.LayerNorm(fc_out_size),
nn.LeakyReLU(0.2),
nn.Linear(fc_out_size, fc_out_size),
)
