def create_decoder(latent, bn_flag):
l1 = Conv2d([latent], [l_dims[2][0]],
l_dims[1][0],
kernel_size=l_dims[2][1:3],
name="dec1",
random_state=random_state)
bn_l1 = BatchNorm2d(l1, bn_flag, name="bn_dec3")
r_l1 = ReLU(bn_l1)
l2 = ConvTranspose2d([r_l1], [l_dims[1][0]],
l_dims[0][0],
kernel_size=l_dims[1][1:3],
name="dec2",
strides=l_dims[1][-1],
border_mode=bpad,
random_state=random_state)
bn_l2 = BatchNorm2d(l2, bn_flag, name="bn_dec2")
r_l2 = ReLU(bn_l2)
l3 = ConvTranspose2d([r_l2], [l_dims[0][0]],
1,
kernel_size=l_dims[0][1:3],
name="dec3",
strides=l_dims[0][-1],
border_mode=bpad,
random_state=random_state)
s_l3 = Sigmoid(l3)
return s_l3
def create_decoder(latent, bn_flag):
l1 = Conv2d([latent], [l_dims[3][0]], l_dims[2][0], kernel_size=l_dims[3][1:3], name="dec1",
random_state=random_state)
bn_l1 = BatchNorm2d(l1, bn_flag, name="bn_dec1")
r_l1 = ReLU(bn_l1)
l2 = ConvTranspose2d([r_l1], [l_dims[2][0]], l_dims[1][0], kernel_size=l_dims[2][1:3], name="dec2",
strides=l_dims[2][-1],
border_mode=dbpad,
random_state=random_state)
bn_l2 = BatchNorm2d(l2, bn_flag, name="bn_dec2")
r_l2 = ReLU(bn_l2)
l3 = ConvTranspose2d([r_l2], [l_dims[1][0]], l_dims[0][0], kernel_size=l_dims[1][1:3], name="dec3",
strides=l_dims[1][-1],
border_mode=dbpad,
random_state=random_state)
bn_l3 = BatchNorm2d(l3, bn_flag, name="bn_dec3")
r_l3 = ReLU(bn_l3)
# hack it and do depth to space
in_chan = l_dims[0][0]
out_chan = 257
kernel_sz = [l_dims[0][1], l_dims[0][2]]
kernel_sz[1] = 1
l4 = ConvTranspose2d([r_l3], [in_chan], out_chan, kernel_size=kernel_sz, name="dec4",
strides=l_dims[0][-1],
border_mode=dbpad,
random_state=random_state)
s_l4 = Sigmoid(l4)
s_l4 = tf.transpose(s_l4, (0, 1, 3, 2))
return s_l4
def create_encoder(inp, bn_flag):
l1 = Conv2d([inp], [1],
l_dims[0][0],
kernel_size=l_dims[0][1:3],
name="enc1",
strides=l_dims[0][-1],
border_mode=bpad,
random_state=random_state)
bn_l1 = BatchNorm2d(l1, bn_flag, name="bn_enc1")
r_l1 = ReLU(bn_l1)
l2 = Conv2d([r_l1], [l_dims[0][0]],
l_dims[1][0],
kernel_size=l_dims[1][1:3],
name="enc2",
strides=l_dims[1][-1],
border_mode=bpad,
random_state=random_state)
bn_l2 = BatchNorm2d(l2, bn_flag, name="bn_enc2")
r_l2 = ReLU(bn_l2)
l3 = Conv2d([r_l2], [l_dims[1][0]],
l_dims[2][0],
kernel_size=l_dims[2][1:3],
name="enc3",
random_state=random_state)
bn_l3 = BatchNorm2d(l3, bn_flag, name="bn_enc3")
return bn_l3
def create_decoder(latent, bn_flag):
l1 = Conv2d([latent], [l_dims[-1][0]],
l_dims[-2][0],
kernel_size=l_dims[-1][1:3],
name="dec1",
random_state=random_state)
bn_l1 = BatchNorm2d(l1, bn_flag, name="bn_dec1")
r_l1 = ReLU(bn_l1)
l2 = ConvTranspose2d([r_l1], [l_dims[-2][0]],
l_dims[-3][0],
kernel_size=l_dims[-2][1:3],
name="dec2",
strides=l_dims[-2][-1],
border_mode=dbpad,
random_state=random_state)
bn_l2 = BatchNorm2d(l2, bn_flag, name="bn_dec2")
r_l2 = ReLU(bn_l2)
l3 = ConvTranspose2d([r_l2], [l_dims[-3][0]],
4 * n_out,
kernel_size=l_dims[-3][1:3],
name="dec3",
strides=l_dims[-3][-1],
border_mode=dbpad,
random_state=random_state)
return tf.reshape(l3, (-1, 1, 48, 4, n_out))
def create_decoder(latent, bn_flag):
l1 = Conv2d([latent], [l_dims[-1][0]],
l_dims[-2][0],
kernel_size=l_dims[-1][1:3],
name="dec1",
random_state=random_state)
bn_l1 = BatchNorm2d(l1, bn_flag, name="bn_dec1")
r_l1 = ReLU(bn_l1)
l2 = ConvTranspose2d([r_l1], [l_dims[-2][0]],
l_dims[-3][0],
kernel_size=l_dims[-2][1:3],
name="dec2",
strides=l_dims[-2][-1],
border_mode=dbpad,
random_state=random_state)
bn_l2 = BatchNorm2d(l2, bn_flag, name="bn_dec2")
r_l2 = ReLU(bn_l2)
l3 = ConvTranspose2d([r_l2], [l_dims[-3][0]],
l_dims[-4][0],
kernel_size=l_dims[-3][1:3],
name="dec3",
strides=l_dims[-3][-1],
border_mode=dbpad,
random_state=random_state)
bn_l3 = BatchNorm2d(l3, bn_flag, name="bn_dec3")
r_l3 = ReLU(bn_l3)
l4 = ConvTranspose2d([r_l3], [l_dims[-4][0]],
l_dims[-5][0],
kernel_size=l_dims[-4][1:3],
name="dec4",
strides=l_dims[-4][-1],
border_mode=dbpad,
random_state=random_state)
bn_l4 = BatchNorm2d(l4, bn_flag, name="bn_dec4")
r_l4 = ReLU(bn_l4)
l5 = ConvTranspose2d([r_l4], [l_dims[-5][0]],
dmol_proj,
kernel_size=l_dims[-5][1:3],
name="dec5",
strides=l_dims[-5][-1],
border_mode=dbpad,
random_state=random_state)
l5_mix, l5_means, l5_lin_scales = DiscreteMixtureOfLogistics(
[l5], [dmol_proj],
n_components=n_components,
name="d_out",
random_state=random_state)
return l5_mix, l5_means, l5_lin_scales
def create_decoder(latent, bn_flag):
l1 = Conv2d([latent], [l_dims[-1][0]],
l_dims[-2][0],
kernel_size=l_dims[-1][1:3],
name="dec1",
random_state=random_state)
bn_l1 = BatchNorm2d(l1, bn_flag, name="bn_dec1")
r_l1 = ReLU(bn_l1)
l2 = ConvTranspose2d([r_l1], [l_dims[-2][0]],
l_dims[-3][0],
kernel_size=l_dims[-2][1:3],
name="dec2",
strides=l_dims[-2][-1],
border_mode=dbpad,
random_state=random_state)
bn_l2 = BatchNorm2d(l2, bn_flag, name="bn_dec2")
r_l2 = ReLU(bn_l2)
l3 = ConvTranspose2d([r_l2], [l_dims[-3][0]],
l_dims[-4][0],
kernel_size=l_dims[-3][1:3],
name="dec3",
strides=l_dims[-3][-1],
border_mode=dbpad,
random_state=random_state)
bn_l3 = BatchNorm2d(l3, bn_flag, name="bn_dec3")
r_l3 = ReLU(bn_l3)
l4 = ConvTranspose2d([r_l3], [l_dims[-4][0]],
l_dims[-5][0],
kernel_size=l_dims[-4][1:3],
name="dec4",
strides=l_dims[-4][-1],
border_mode=dbpad,
random_state=random_state)
bn_l4 = BatchNorm2d(l4, bn_flag, name="bn_dec4")
r_l4 = ReLU(bn_l4)
l5 = ConvTranspose2d([r_l4], [l_dims[-5][0]],
1,
kernel_size=l_dims[-5][1:3],
name="dec5",
strides=l_dims[-5][-1],
border_mode=dbpad,
random_state=random_state)
#s_l5 = Sigmoid(l5)
t_l5 = Tanh(l5)
return t_l5
def create_pixel_cnn(inp, lbl, cond):
e_inp, emb = Embedding(inp, 256, n_channels, random_state=random_state, name="inp_emb")
c_inp, emb = Embedding(cond, 256, n_channels, random_state=random_state, name="cond_emb")
l1_v, l1_h = GatedMaskedConv2d([e_inp], [n_channels], [e_inp], [n_channels],
n_channels,
residual=False,
conditioning_class_input=lbl,
conditioning_num_classes=n_labels,
conditioning_spatial_map=c_inp,
kernel_size=kernel_size0, name="pcnn0",
mask_type="img_A",
random_state=random_state)
o_v = l1_v
o_h = l1_h
for i in range(n_layers - 1):
t_v, t_h = GatedMaskedConv2d([o_v], [n_channels], [o_h], [n_channels],
n_channels,
conditioning_class_input=lbl,
conditioning_num_classes=n_labels,
conditioning_spatial_map=c_inp,
kernel_size=kernel_size1, name="pcnn{}".format(i + 1),
mask_type="img_B",
random_state=random_state)
o_v = t_v
o_h = t_h
cleanup = Conv2d([o_h], [n_channels], n_channels, kernel_size=(1, 1),
name="conv_c",
random_state=random_state)
r_p = ReLU(cleanup)
out = Conv2d([r_p], [n_channels], 256, kernel_size=(1, 1),
name="conv_o",
random_state=random_state)
#s_out = Softmax(out)
return out#s_out
def create_encoder(inp, bn_flag):
e_inps = []
for ci in range(4):
e_inp, emb = Embedding(inp[..., ci][..., None],
n_out,
inp_emb_dim,
random_state=random_state,
name="inp_emb_{}".format(ci))
e_inps.append(e_inp)
e_inp = tf.concat(e_inps, axis=-1)
l1 = Conv2d([e_inp], [4 * inp_emb_dim],
l_dims[0][0],
kernel_size=l_dims[0][1:3],
name="enc1",
strides=l_dims[0][-1],
border_mode=ebpad,
random_state=random_state)
bn_l1 = BatchNorm2d(l1, bn_flag, name="bn_enc1")
r_l1 = ReLU(bn_l1)
l2 = Conv2d([r_l1], [l_dims[0][0]],
l_dims[1][0],
kernel_size=l_dims[1][1:3],
name="enc2",
strides=l_dims[1][-1],
border_mode=ebpad,
random_state=random_state)
bn_l2 = BatchNorm2d(l2, bn_flag, name="bn_enc2")
r_l2 = ReLU(bn_l2)
l3 = Conv2d([r_l2], [l_dims[1][0]],
l_dims[2][0],
kernel_size=l_dims[2][1:3],
name="enc3",
random_state=random_state)
bn_l3 = BatchNorm2d(l3, bn_flag, name="bn_enc3")
return bn_l3
def create_model(inp_tm1, inp_t, h1_init, c1_init, h1_q_init, c1_q_init):
p_tm1 = Linear([inp_tm1], [1],
n_hid,
random_state=random_state,
name="proj_in",
init=forward_init)
def step(x_t, h1_tm1, c1_tm1, h1_q_tm1, c1_q_tm1):
output, s = LSTMCell([x_t, h1_q_tm1], [n_hid, n_hid],
h1_tm1,
c1_tm1,
n_hid,
random_state=random_state,
name="rnn1",
init=rnn_init)
h1_t = s[0]
c1_t = s[1]
output, s = LSTMCell([h1_t], [n_hid],
h1_q_tm1,
c1_q_tm1,
n_hid,
random_state=random_state,
name="rnn1_q",
init=rnn_init)
h1_cq_t = s[0]
c1_q_t = s[1]
h1_q_t, h1_i_t, h1_nst_q_t, h1_emb = VqEmbedding(
h1_cq_t, n_hid, n_emb, random_state=random_state, name="h1_vq_emb")
output_q_t, output_i_t, output_nst_q_t, output_emb = VqEmbedding(
output, n_hid, n_emb, random_state=random_state, name="out_vq_emb")
# not great
h1_i_t = tf.cast(h1_i_t, tf.float32)
output_i_t = tf.cast(h1_i_t, tf.float32)
lf_output = Bilinear(h1_q_t,
n_hid,
output_emb,
n_hid,
random_state=random_state,
name="out_mix",
init=forward_init)
rf_output = Bilinear(output_q_t,
n_hid,
h1_emb,
n_hid,
random_state=random_state,
name="h_mix",
init=forward_init)
f_output = Linear([lf_output, rf_output], [n_emb, n_emb],
n_hid,
random_state=random_state,
name="out_f",
init=forward_init)
# r[0]
rets = [f_output]
# r[1:3]
rets += [h1_t, c1_t]
# r[3:9]
rets += [h1_q_t, c1_q_t, h1_nst_q_t, h1_cq_t, h1_i_t, h1_emb]
# r[9:]
rets += [output_q_t, output_nst_q_t, output, output_i_t, output_emb]
return rets
outputs_info = [
None, h1_init, c1_init, h1_q_init, c1_q_init, None, None, None, None,
None, None, None, None, None
]
r = scan(step, [p_tm1], outputs_info)
out = r[0]
hiddens = r[1]
cells = r[2]
q_hiddens = r[3]
q_cells = r[4]
q_nst_hiddens = r[5]
q_nvq_hiddens = r[6]
i_hiddens = r[7]
emb_hiddens = r[8]
q_out = r[9]
q_nst_out = r[10]
q_nvq_out = r[11]
i_out = r[12]
emb_out = r[13]
l1 = Linear([out, q_hiddens], [n_hid, n_hid],
n_hid,
random_state=random_state,
name="l1",
init=forward_init)
r_l1 = ReLU(l1)
pred = Linear([r_l1], [n_hid],
1,
random_state=random_state,
name="out",
init=forward_init)
outs_names = [
"pred", "hiddens", "cells", "q_hiddens", "q_cells", "q_nst_hiddens",
"q_nvq_hiddens", "i_hiddens", "emb_hiddens", "q_out", "q_nst_out",
"q_nvq_out", "i_out", "emb_out"
]
outs_tf = [eval(name) for name in outs_names]
c = namedtuple("Core", outs_names)(*outs_tf)
return c
def create_model(inp_tm1, h1_q_init):
def step(x_t, h1_tm1):
output, s = GRUCell([x_t], [1],
h1_tm1,
n_hid,
random_state=random_state,
name="rnn1",
init=rnn_init)
h1_cq_t = s[0]
"""
output, s = LSTMCell([h1_t], [n_hid], h1_q_tm1, c1_q_tm1, n_hid,
random_state=random_state,
name="rnn1_q", init=rnn_init)
h1_cq_t = s[0]
c1_q_t = s[1]
"""
qhs = []
ihs = []
nst_qhs = []
embs = []
for i in list(range(n_split)):
e_div = int(n_hid / n_split)
h1_q_t, h1_i_t, h1_nst_q_t, h1_emb = VqEmbedding(
h1_cq_t[:, i * e_div:(i + 1) * e_div],
e_div,
n_emb,
random_state=random_state,
# shared space?
name="h1_vq_emb")
#name="h1_{}_vq_emb".format(i))
qhs.append(h1_q_t)
ihs.append(h1_i_t[:, None])
nst_qhs.append(h1_nst_q_t)
embs.append(h1_emb)
h1_q_t = tf.concat(qhs, axis=-1)
h1_nst_q_t = tf.concat(nst_qhs, axis=-1)
h1_i_t = tf.concat(ihs, axis=-1)
# not great
h1_i_t = tf.cast(h1_i_t, tf.float32)
return output, h1_q_t, h1_nst_q_t, h1_cq_t, h1_i_t
r = scan(step, [inp_tm1], [None, h1_q_init, None, None, None])
out = r[0]
q_hiddens = r[1]
q_nst_hiddens = r[2]
q_nvq_hiddens = r[3]
i_hiddens = r[4]
l1 = Linear([out], [n_hid],
n_hid,
random_state=random_state,
name="l1",
init=forward_init)
r_l1 = ReLU(l1)
pred = Linear([r_l1], [n_hid],
1,
random_state=random_state,
name="out",
init=forward_init)
return pred, q_hiddens, q_nst_hiddens, q_nvq_hiddens, i_hiddens