keras打印loss对权重的导数

Notes

怀疑模型梯度爆炸，想打印模型 loss 对各权重的导数看看。如果如果fit来训练的话，可以用keras.callbacks.TensorBoard实现。
但此次使用train_on_batch来训练的，用K.gradients和K.function实现。

Codes

• 以一份 VAE 代码为例

# -*- coding: utf8 -*-
import keras
from keras.models import Model
from keras.layers import Input, Lambda, Conv2D, MaxPooling2D, Flatten, Dense, Reshape
from keras.losses import binary_crossentropy
from keras.datasets import mnist, fashion_mnist
import keras.backend as K
from scipy.stats import norm
import numpy as np
import matplotlib.pyplot as plt

BATCH = 128
N_CLASS = 10
EPOCH = 5
IN_DIM = 28 * 28
H_DIM = 128
Z_DIM = 2

(x_train, y_train), (x_test, y_test) = fashion_mnist.load_data()
x_train = x_train.reshape(len(x_train), -1).astype('float32') / 255.
x_test = x_test.reshape(len(x_test), -1).astype('float32') / 255.


def sampleing(args):
    """reparameterize"""
    mu, logvar = args
    eps = K.random_normal([K.shape(mu)[0], Z_DIM], mean=0.0, stddev=1.0)
    return mu + eps * K.exp(logvar / 2.)


# encode
x_in = Input([IN_DIM])
h = Dense(H_DIM, activation='relu')(x_in)
z_mu = Dense(Z_DIM)(h)  # mean，不用激活
z_logvar = Dense(Z_DIM)(h)  # log variance，不用激活
z = Lambda(sampleing, output_shape=[Z_DIM])([z_mu, z_logvar])  # 只能有一个参数
encoder = Model(x_in, [z_mu, z_logvar, z], name='encoder')

# decode
z_in = Input([Z_DIM])
h_hat = Dense(H_DIM, activation='relu')(z_in)
x_hat = Dense(IN_DIM, activation='sigmoid')(h_hat)
decoder = Model(z_in, x_hat, name='decoder')

# VAE
x_in = Input([IN_DIM])
x = x_in
z_mu, z_logvar, z = encoder(x)
x = decoder(z)
out = x
vae = Model(x_in, [out, out], name='vae')


# loss_kl = 0.5 * K.sum(K.square(z_mu) + K.exp(z_logvar) - 1. - z_logvar, axis=1)
# loss_recon = binary_crossentropy(K.reshape(vae_in, [-1, IN_DIM]), vae_out) * IN_DIM
# loss_vae = K.mean(loss_kl + loss_recon)

def loss_kl(y_true, y_pred):
    return 0.5 * K.sum(K.square(z_mu) + K.exp(z_logvar) - 1. - z_logvar, axis=1)


# vae.add_loss(loss_vae)
vae.compile(optimizer='rmsprop',
            loss=[loss_kl, 'binary_crossentropy'],
            loss_weights=[1, IN_DIM])
vae.summary()

# 获取模型权重 variable
w = vae.trainable_weights
print(w)

# 打印 KL 对权重的导数
# KL 要是 Tensor，不能是上面的函数 `loss_kl`
grad = K.gradients(0.5 * K.sum(K.square(z_mu) + K.exp(z_logvar) - 1. - z_logvar, axis=1),
                   w)
print(grad)  # 有些是 None 的
grad = grad[grad is not None]  # 去掉 None，不然报错

# 打印梯度的函数
# K.function 的输入和输出必要是 list！就算只有一个
show_grad = K.function([vae.input], [grad])


# vae.fit(x_train, # y_train,  # 不能传 y_train
#         batch_size=BATCH,
#         epochs=EPOCH,
#         verbose=1,
#         validation_data=(x_test, None))

''' 以 train_on_batch 方式训练 '''
for epoch in range(EPOCH):
    for b in range(x_train.shape[0] // BATCH):
        idx = np.random.choice(x_train.shape[0], BATCH)
        x = x_train[idx]
        l = vae.train_on_batch([x], [x, x])

    # 计算梯度
    gd = show_grad([x])
    # 打印梯度
    print(gd)

# show manifold
PIXEL = 28
N_PICT = 30
grid_x = norm.ppf(np.linspace(0.05, 0.95, N_PICT))
grid_y = grid_x

figure = np.zeros([N_PICT * PIXEL, N_PICT * PIXEL])
for i, xi in enumerate(grid_x):
    for j, yj in enumerate(grid_y):
        noise = np.array([[xi, yj]])  # 必须秩为 2，两层中括号
        x_gen = decoder.predict(noise)
        # print('x_gen shape:', x_gen.shape)
        x_gen = x_gen[0].reshape([PIXEL, PIXEL])
        figure[i * PIXEL: (i+1) * PIXEL,
               j * PIXEL: (j+1) * PIXEL] = x_gen

fig = plt.figure(figsize=(10, 10))
plt.imshow(figure, cmap='Greys_r')
fig.savefig('./variational_autoencoder.png')
plt.show()

References

1. keras.callbacks.TensorBoard
2. keras+TensorBoard实现训练可视化
3. 获取keras模型的weights和bias
4. “让Keras更酷一些！”：中间变量、权重滑动和安全生成器
5. Keras使用K.function抽取中间层报错： TypeError: inputs to a TensorFlow backend function should be a list or t
6. keras/examples/variational_autoencoder.py