def plot_latent_units(mean, std, w):
# plot the latents and its weights
fig = plt.figure(figsize=(6, 4), dpi=200)
ax = plt.gca()
l1 = ax.plot(mean,
label='mean',
linewidth=1.0,
marker='o',
markersize=6,
color='r',
alpha=0.5)
l2 = ax.plot(std,
label='std',
linewidth=1.0,
marker='o',
markersize=6,
color='g',
alpha=0.5)
ax1 = ax.twinx()
l3 = ax1.plot(w,
label='weight',
linewidth=1.0,
linestyle='--',
marker='s',
markersize=6,
color='b',
alpha=0.5)
lines = l1 + l2 + l3
labs = [l.get_label() for l in lines]
ax.grid(True)
ax.legend(lines, labs)
return vs.plot_to_image(fig)
def to_image(X, grids):
if X.shape[-1] == 1: # grayscale image
X = np.squeeze(X, axis=-1)
else: # color image
X = np.transpose(X, (0, 3, 1, 2))
nrows, ncols = grids
fig = vs.plot_figure(nrows=nrows, ncols=ncols, dpi=100)
vs.plot_images(X, grids=grids)
image = vs.plot_to_image(fig)
return image
def llk_pixels(model: VariationalModel, valid_ds: tf.data.Dataset):
llk = []
for x, y in valid_ds.take(5):
px, _ = _call(model, x, y, decode=True)
px = as_tuple(px)[0]
if hasattr(px, 'distribution'):
px = px.distribution
if isinstance(px, Bernoulli):
px = Bernoulli(logits=px.logits)
elif isinstance(px, Normal):
px = Normal(loc=px.loc, scale=px.scale)
elif isinstance(px, QuantizedLogistic):
px = QuantizedLogistic(loc=px.loc,
scale=px.scale,
low=px.low,
high=px.high,
inputs_domain=px.inputs_domain,
reinterpreted_batch_ndims=None)
else:
return # nothing to do
llk.append(px.log_prob(x))
# average over all channels
llk_image = tf.reduce_mean(tf.reduce_mean(tf.concat(llk, 0), axis=0),
axis=-1)
llk = tf.reshape(llk_image, -1)
tf.summary.histogram('valid/llk_pixels', llk, step=model.step)
# show the image heatmap of llk pixels
fig = plt.figure(figsize=(3, 3))
ax = plt.gca()
im = ax.pcolormesh(llk_image.numpy(),
cmap='Spectral',
vmin=np.min(llk),
vmax=np.max(llk))
ax.axis('off')
ax.margins(0.)
# color bar
ticks = np.linspace(np.min(llk), np.max(llk), 5)
cbar = plt.colorbar(im, ax=ax, fraction=0.04, pad=0.02, ticks=ticks)
cbar.ax.set_yticklabels([f'{i:.2f}' for i in ticks])
cbar.ax.tick_params(labelsize=6)
plt.tight_layout()
tf.summary.image('llk_heatmap', vs.plot_to_image(fig, dpi=100))
def callback(vae: vi.VariationalAutoencoder, x: np.ndarray, y: np.ndarray):
trainer = get_current_trainer()
px, qz = [], []
X_i = []
for x_i in tf.data.Dataset.from_tensor_slices(x).batch(64):
_ = vae(x_i, training=False)
px.append(_[0])
qz.append(_[1])
X_i.append(x_i)
# llk
llk_test = tf.reduce_mean(
tf.concat([p.log_prob(x_i) for p, x_i in zip(px, X_i)], axis=0))
# latents
qz_mean = tf.reduce_mean(tf.concat([q.mean() for q in qz], axis=0), axis=0)
qz_std = tf.reduce_mean(tf.concat([q.stddev() for q in qz], axis=0),
axis=0)
w = tf.reduce_sum(vae.decoder.trainable_variables[0], axis=1)
# plot the latents and its weights
fig = plt.figure(figsize=(6, 4), dpi=200)
ax = plt.gca()
l1 = ax.plot(qz_mean,
label='mean',
linewidth=1.0,
linestyle='--',
marker='o',
markersize=4,
color='r',
alpha=0.5)
l2 = ax.plot(qz_std,
label='std',
linewidth=1.0,
linestyle='--',
marker='o',
markersize=4,
color='g',
alpha=0.5)
ax1 = ax.twinx()
l3 = ax1.plot(w,
label='weight',
linewidth=1.0,
linestyle='--',
marker='o',
markersize=4,
color='b',
alpha=0.5)
lines = l1 + l2 + l3
labs = [l.get_label() for l in lines]
ax.grid(True)
ax.legend(lines, labs)
img_qz = vs.plot_to_image(fig)
# reconstruction
img = px[10].mean().numpy()
if img.shape[-1] == 1:
img = np.squeeze(img, axis=-1)
fig = plt.figure(figsize=(8, 8), dpi=120)
vs.plot_images(img, grids=(8, 8))
img_reconstructed = vs.plot_to_image(fig)
# latents traverse
# TODO
return dict(llk_test=llk_test,
qz_mean=qz_mean,
qz_std=qz_std,
w_decoder=w,
reconstructed=img_reconstructed,
latents=img_qz)
def latent_units(model: VariationalModel, valid_ds: tf.data.Dataset):
weights = model.weights
Qz = []
Pz = []
for x, y in valid_ds.take(10):
_call(model, x, y, decode=True)
qz, pz = model.get_latents(return_prior=True)
Qz.append(as_tuple(qz))
Pz.append(as_tuple(pz))
n_latents = len(Qz[0])
for i in range(n_latents):
qz: Sequence[Distribution] = [q[i] for q in Qz]
pz: Sequence[Distribution] = [p[i] for p in Pz]
# tracking kl
kld = []
for q, p in zip(qz, pz):
q = Normal(loc=q.mean(), scale=q.stddev())
z = q.sample()
if isinstance(p, Vamprior):
C = p.C
p = p.distribution # [n_components, zdim]
p = Normal(loc=p.mean(), scale=p.stddev())
kld.append(
q.log_prob(z) - (tf.reduce_logsumexp(
p.log_prob(tf.expand_dims(z, 1)), 1) - C))
else:
p = Normal(loc=p.mean(), scale=p.stddev())
kld.append(q.log_prob(z) - p.log_prob(z))
kld = tf.reshape(tf.reduce_mean(tf.concat(kld, 0), axis=0),
-1).numpy()
# mean and stddev
mean = tf.reduce_mean(tf.concat([d.mean() for d in qz], axis=0), 0)
stddev = tf.reduce_mean(
tf.concat([d.stddev() for d in qz], axis=0), 0)
mean = tf.reshape(mean, -1).numpy()
stddev = tf.reshape(stddev, -1).numpy()
zdim = mean.shape[0]
# the figure
plt.figure(figsize=(12, 5), dpi=50)
lines = []
ids = np.argsort(stddev)
styles = dict(marker='o', markersize=2, linewidth=0, alpha=0.8)
lines += plt.plot(mean[ids], label='mean', color='r', **styles)
lines += plt.plot(stddev[ids], label='stddev', color='b', **styles)
plt.grid(False)
# show weights if exists
plt.twinx()
lines += plt.plot(kld[ids],
label='KL(q|p)',
linestyle='--',
linewidth=1.0,
alpha=0.6)
for w in weights:
name = w.name
if w.shape.rank > 0 and w.shape[
0] == zdim and '/kernel' in name:
w = tf.linalg.norm(tf.reshape(w, (w.shape[0], -1)),
axis=1).numpy()
lines += plt.plot(w[ids],
label=name.split(':')[0],
linestyle='--',
alpha=0.6)
plt.grid(False)
plt.legend(lines, [ln.get_label() for ln in lines], fontsize=6)
# save summary
tf.summary.image(f'z{i}',
vs.plot_to_image(plt.gcf()),
step=model.step)
tf.summary.histogram(f'z{i}/mean', mean, step=model.step)
tf.summary.histogram(f'z{i}/stddev', stddev, step=model.step)
tf.summary.histogram(f'z{i}/kld', kld, step=model.step)
def callback():
trainer = get_current_trainer()
x, y = x_test[:1000], y_test[:1000]
px, qz = vae(x, training=False)
# latents
qz_mean = tf.reduce_mean(qz.mean(), axis=0)
qz_std = tf.reduce_mean(qz.stddev(), axis=0)
w = tf.reduce_sum(decoder.trainable_variables[0], axis=(0, 1, 2))
# plot the latents and its weights
fig = plt.figure(figsize=(6, 4), dpi=200)
ax = plt.gca()
l1 = ax.plot(qz_mean,
label='mean',
linewidth=1.0,
linestyle='--',
marker='o',
markersize=4,
color='r',
alpha=0.5)
l2 = ax.plot(qz_std,
label='std',
linewidth=1.0,
linestyle='--',
marker='o',
markersize=4,
color='g',
alpha=0.5)
ax1 = ax.twinx()
l3 = ax1.plot(w,
label='weight',
linewidth=1.0,
linestyle='--',
marker='o',
markersize=4,
color='b',
alpha=0.5)
lines = l1 + l2 + l3
labs = [l.get_label() for l in lines]
ax.grid(True)
ax.legend(lines, labs)
img_qz = vs.plot_to_image(fig)
# reconstruction
fig = plt.figure(figsize=(5, 5), dpi=120)
vs.plot_images(np.squeeze(px.mean().numpy()[:25], axis=-1),
grids=(5, 5))
img_res = vs.plot_to_image(fig)
# latents
fig = plt.figure(figsize=(5, 5), dpi=200)
z = fast_umap(qz.mean().numpy())
vs.plot_scatter(z, color=y, size=12.0, alpha=0.4)
img_umap = vs.plot_to_image(fig)
# gradients
grads = [(k, v) for k, v in trainer.last_train_metrics.items()
if '_grad/' in k]
encoder_grad = sum(v for k, v in grads if 'Encoder' in k)
decoder_grad = sum(v for k, v in grads if 'Decoder' in k)
return dict(reconstruct=img_res,
umap=img_umap,
latents=img_qz,
qz_mean=qz_mean,
qz_std=qz_std,
w_decoder=w,
llk_test=tf.reduce_mean(px.log_prob(x)),
encoder_grad=encoder_grad,
decoder_grad=decoder_grad)