TF2下变分自编码的N种写法
在开篇之前，请允许我吐槽几段文字，发泄一下TF的不便之处。如果对这部分内容不敢兴趣请直接看正文内容。
【吐槽部分】：
在TF升级到2.x之后，带给读者更多的编码方式，同时也带给读者更多的坑。使得本来比较难用的TensorFlow变得又灵活，又难用。这使得好多TensorFlow下的深度用户，一夜之间对该框架感到陌生。
能把一个通用框架改到让原有用户都陌生的地步，这也是需要功底的。在1.x时代，一个模型的只有一种写法，规则晦涩很容易出错。在2.x时代，一个模型会变成有N种写法，而且每种写法的规则更加晦涩，写起模型来，出错率成指数增长。这使得开发人员一般会用30%的时间来实现逻辑，70%的时间来处理各种框架运行时所遇到的软件问题。大大降低了开发效率。
在《深度学习之TensorFlow：工程化项目实战》一书中，介绍了TF框架中不下于10种的子开发框架。每种都自立门户，互不兼容。随着TF2.x的到来，砍掉了好多子框架。这使得原本小范围内互不兼容的场面变成整体版本间的绝对不兼容，可以说是将不兼容属性发挥到了极致。
不过透过TF2.x的自杀式改革背后，可以看出，其希望扭转这一困境的决心。在TF2.x中，主推了2个子框架，keras与原生的动态图框架。大概这将会是TF2.x未来的使用趋势。
然而，即便是这两个子框架，自由组合起来，也可以实现n中开发方式。对用户来说，还是一样的灵活、坑多。本文就TF2.x在这两个框架下的开发，做一个系统的介绍。我们尽量不发散太多的开发方法。只针对最主流、最常用的开发方式进行介绍。也希望读者可以真正精通掌握其中的一个开发方法，至少在开发过程中，可以少一些调试框架的时间。
【正文部分】：
在《深度学习之TensorFlow：入门、原理与进阶实战》一书中，第10章介绍过变分自编码以及其在TF1.x下静态图模式的代码实现。该模型的结构相对来讲较为奇特，选用其作为例子讲解，可以触碰到更多开发中遇到的特殊情况。
为了将主流的TF2.x开发模式讲透，这里选用了与书中一样的模型和MNIST数据集。使用tf.Keras接口进行搭建模型，使用keras和动态图两种方式进行训练模型。
在学习本文之前，请先熟悉一下书中的变分自编码介绍。我们以前发表过的一篇文章<TensorFlow 2.0中实现自动编码器>
1 基础的Keras写法
先来看看最基础的keras写法
1.1 模型结构
解码器与编码器的结构代码如下:
batch_size = 100
original_dim = 784 #28*28
latent_dim = 2
intermediate_dim = 256
nb_epoch = 50
class Encoder(tf.keras.Model): # 编码器
def __init__(self ,intermediate_dim,latent_dim, **kwargs):
super(Encoder, self).__init__(**kwargs)
self.hidden_layer = Dense(units=intermediate_dim, activation=tf.nn.relu)
self.z_mean = Dense(units=latent_dim)
self.z_log_var = Dense(units=latent_dim)
def call(self, x):
activation = self.hidden_layer(x)
z_mean = self.z_mean(activation)
z_log_var= self.z_log_var(activation)
return z_mean,z_log_var
class Decoder(tf.keras.Model): # 解码器
def __init__(self ,intermediate_dim,original_dim,**kwargs):
super(Decoder, self).__init__(**kwargs)
self.hidden_layer = Dense(units=intermediate_dim, activation=tf.nn.relu)
self.output_layer = Dense(units=original_dim, activation='sigmoid')
def call(self, z):
activation = self.hidden_layer(z)
output_layer = self.output_layer(activation)
return output_layer
def samplingfun(z_mean, z_log_var): #采样函数
epsilon = K.random_normal(shape=(K.shape(z_mean)[0], latent_dim),
mean=0., stddev=1.0)
return z_mean + K.exp(z_log_var / 2) * epsilon
def sampling(args):
z_mean, z_log_var = args
return samplingfun(z_mean, z_log_var)
模型很简单只有全连接神经网络。
1.2 组合模型
定义采样器,并将编码器和解码器组合起来,形成变分自编码模型.
encoder = Encoder(intermediate_dim,latent_dim)
decoder = Decoder(intermediate_dim,original_dim)
inputs = Input(batch_shape=(batch_size, original_dim))
z_mean,z_log_var = encoder(inputs)
z= samplingfun(z_mean, z_log_var)
y_pred = decoder(z)
autoencoder = Model(inputs, y_pred, name='autoencoder')
autoencoder.summary()定义Input是Keras标准的使用技巧.详细介绍可以参考《深度学习之TensorFlow：工程化项目实战》一书第6章
1.3 坑1 :keras自定义模型的默认输入
如果在TF1.x中代码第1.2小节第7行会有问题,它是一个函数不能充当一个层.必须将其封装成层才行.
在TF2.x中,代码第1.2小节第7行是没问题的.但是也不正规,如果运行两次(将第1.2小节第7行代码重复一下),则会报以下错误:
这是个很难查出原因的错误.一个隐形的坑.所以最好还是用Lambda封装成一个层来使用,封装后,运行2次将不会报错.
【坑】:在使用Lambda时,被封装的函数必须只能有一个参数. 比较下面两种写法,
正确的:
z = Lambda(sampling, output_shape=(latent_dim,))([z_mean, z_log_var])错误的:
z = Lambda(samplingfun, output_shape=(latent_dim,))(z_mean, z_log_var)错误的写法,会报如下错误:
大概意思是只向samplingfun传入了一个参数. Samplingfun没有收到z_log_var
1.4 损失函数和编译模型
用二进制交叉熵做重建损失，在配合KL散度损失，具体代码如下：
def vae_loss(x, x_decoded_mean):
xent_loss = original_dim * metrics.binary_crossentropy(x, x_decoded_mean)
kl_loss = - 0.5 * K.sum(1 + z_log_var - K.square(z_mean) - K.exp(z_log_var), axis=-1)
return xent_loss + kl_loss
#编译模型
autoencoder.compile(optimizer='rmsprop', loss=vae_loss)其中重建损失也可以用MSE来代替。例如53行可以写成：
xent_loss = 0.5 * K.sum(K.square(x_decoded_mean - x), axis=-1)模型必须编译才能使用
1.5 载入数据集
载入MNIST数据集
(x_train, y_train), (x_test, y_test) = mnist.load_data(path='mnist.pkl.gz')
x_train = x_train.astype('float32') / 255.
x_test = x_test.astype('float32') / 255.
x_train = x_train.reshape((len(x_train), np.prod(x_train.shape[1:])))
x_test = x_test.reshape((len(x_test), np.prod(x_test.shape[1:])))1.6 训练模型
使用keras原生的fit来训练模型.
autoencoder.fit(x_train, x_train,
shuffle=True,
epochs=5,
verbose=2,
batch_size=batch_size,
validation_data=(x_test, x_test))fit是keras中非常常用的训练框架.详细介绍可以参考书中内容,这里不再展开.
1.7 使用模型
使用Model可以将任意张量组成模型下面第1行组成了一个输入是inputs输出是z_mean的模型,用该模型输出数据集中的解码均值,并显示出来.
modencoder = Model(inputs, z_mean)
print(modencoder.layers[1]) #Encoder
x_test_encoded = modencoder.predict(x_test, batch_size=batch_size)
plt.figure(figsize=(6, 6))
plt.scatter(x_test_encoded[:, 0], x_test_encoded[:, 1], c=y_test)
plt.colorbar()
plt.show()输出如下:
2 无监督训练中,没有标签的代码如何编写
在1中,介绍的训练方式是典形的有标签训练.即,在训练模型时,输入了2个样本,都是x_train.(代码第1.6小节第1行)
还可以稍加修改,在fit时不传入标签y,这样可以提升运算效率.具体做法如下:
2.1 修改组合模型
将1.2中组合模型和1.4的损失函数代码修改如下:
encoder = Encoder(intermediate_dim,latent_dim)
decoder = Decoder(intermediate_dim,original_dim)
inputs = Input(batch_shape=(batch_size, original_dim))
z_mean,z_log_var = encoder(inputs)
z = Lambda(sampling, output_shape=(latent_dim,))([z_mean, z_log_var])
y_pred = decoder(z)
autoencoder = Model(inputs, y_pred, name='autoencoder')
# 重构loss
xent_loss = original_dim * metrics.binary_crossentropy(inputs, y_pred)
# KL loss
kl_loss = - 0.5 * K.sum(1 + z_log_var -
K.square(z_mean) - K.exp(z_log_var), axis=-1)
vae_loss = K.mean(xent_loss + kl_loss)
autoencoder.add_loss(vae_loss)
autoencoder.compile(optimizer='rmsprop')可以看到,直接将1.4节的损失函数展开,生成了张量vae_loss.同时,利用模型的add_loss方法,将张量损失加入进去.
在编译模型时,可以不需要再指定损失了.
2.2 坑2:向模型中加入损失张量
最常见的坑,就是使用1.4节的方法,将张量损失编译到模型里.写法如下:
autoencoder.compile(optimizer='rmsprop', loss=vae_loss)这时,报的错误绝对会使你一脸懵.错误是这样的:
【坑】:所以,记住张量损失一定要用模型的add_loss方法进行添加
2.3 坑3: 模型与训练不匹配
代码改到这里,并没有完事.因为我们将模型的loss计算中的标签输入去掉了.而fit的时候,还会输入标签.这时如果直接运行会报如下错误:
同样,这个错误也很难看出问题所在.直接修改fit函数,将输入标签去掉即可.代码1.6小节改成如下:
autoencoder.fit(x_train, validation_split=0.05, epochs=5, batch_size=batch_size)再次运行,即可通过.
3 张量损失封装成损失函数
其实2节所介绍的方法,也可以再次封装成损失函数来进行执行.具体做法如下.
3.1 将张量损失封装成函数
在2.1小节代码后面添加如下代码:
def vae_lossfun(x, loss):
return loss
lossautoencoder = Model(inputs, vae_loss, name='lossautoencoder')
lossautoencoder.compile(optimizer='rmsprop', loss=vae_lossfun)该代码,重新又建立一个模型lossautoencoder,该模型的输出就是张量损失vae_loss.同时又建立一个损失函数,在输入损失时,将模型的输出透传出来即可.
在训练时可以直接使用1.6小节的fit方法即可
3.2 坑4:损失函数的参数固定
一定要注意,损失函数的参数是固定的(第一个是标签,第二个是预测值).如果将vae_lossfun,的参数改变,
def vae_lossfun(loss):
运行时,将会出现错误.
当然,这个也可以自定义.再未来的新书<机器视觉之TensorFlow2:入门原理与应用实战>里,会更加全面的展开介绍.
3.3 总结
这种模型的输出就是损失值的情况非常常见.例如最大化互信息模型(DIM)就是这种.2节和3节提供了2个这种模型的训练方法,都可以使用.
4 使用动态图训练
前面的1,2,3节都是使用keras的方式来训练模型.这种方法看是方便,但不适合模型的调试环节.尤其当训练种出现了None,更是一头雾水.虽然keras有单步训练的方式,但是仍不够灵活,为了适应训练过程中,各种情况的调试,最好还是使用底层的动态图训练模型.
4.1 修改训练方式
将1.6的代码改成如下:
optimizer = Adam(lr=0.001)#定义优化器
training_dataset = tf.data.Dataset.from_tensor_slices( #定义数据集
x_train).batch(batch_size)
nb_epoch = 5
for epoch in range(nb_epoch): # 按照指定迭代次数进行训练
for dataone in training_dataset: # 遍历数据集
img = np.reshape(dataone, (batch_size, -1))
with tf.GradientTape() as tape:
z_mean,z_log_var = encoder(img)
z= samplingfun(z_mean, z_log_var) #ok
x_decoded_mean = decoder(z)
xent_loss = K.sum(K.square(x_decoded_mean - img), axis=-1)
kl_loss = - 0.5 * K.sum(1 + z_log_var - K.square(z_mean) - K.exp(z_log_var), axis=-1)
thisloss = K.mean(xent_loss)*0.5+K.mean(kl_loss)
gradients = tape.gradient(thisloss,
encoder.trainable_variables+decoder.trainable_variables)
gradient_variables = zip(gradients,
encoder.trainable_variables+decoder.trainable_variables)
optimizer.apply_gradients(gradient_variables)在动态图中训练需要自己定义数据集.还好,tf中数据集接口比较方便.一行代码就可以搞定. 《深度学习之TensorFlow：工程化项目实战》一书中数据集的介绍花了很大篇幅,相信读者这部分本领应该可以过关.
【坑】:即便是tf2.x其数据集接口的处理函数仍然是静态图,这个给调试过程带来很大不便.解决方法就是将其内部的调用再转成动态图.具体见书中介绍,这里不再展开.
动态图的代码很有流程化,读者只需要按顺序一步一步的做皆可.在倒数后两行需要注意,得将要训练得模型权重全部放进去.
4.2 技巧:任意提取模型
使用动态图训练好的模型,本质是改变实例化模型类的对象.所以也可以再使用keras.model方法,将其任意组成子模型.
这是tf框架非常赞的地方.它可以非常方便的将子模型提取出来,并通过权重的载入载出方法将模型保存和加载,例如:
modeENCODER.save_weights('my_modeldimvae.h5')
modeENCODER.load_weights('my_modeldimvae.h5')
这种方法可以非常方便的进行模型工程化部署.
具体例子,见配套的源码文件
5 以类的方式封装模型损失函数
为了代码工整,还可以将模型的整个过程封装起来,直接输出损失函数.
5.1 封装损失函数
将整个流程封装起来,具体如下:
class VAE(tf.keras.Model): # 提取图片特征
def __init__(self ,intermediate_dim,original_dim,latent_dim,**kwargs):
super(VAE, self).__init__(**kwargs)
self.encoder = Encoder(intermediate_dim,latent_dim)
self.decoder = Decoder(intermediate_dim,original_dim)
def call(self, x):
z_mean,z_log_var = self.encoder(x)
z = sampling( (z_mean,z_log_var) )
y_pred = self.decoder(z)
xent_loss = original_dim * metrics.binary_crossentropy(x, y_pred) #ok
kl_loss = - 0.5 * K.sum(1 + z_log_var - K.square(z_mean) - K.exp(z_log_var), axis=-1)
loss = xent_loss + kl_loss
return loss该类的训练方法可以参考2,3节.具体可以参考配套的源码.
5.2 更合理的类封装模式
真正使用是,常常会将特征提取部分单独分开,作为一个类.这样利于扩展.令变分自编码功能方面的部分单独成一个类只完成变分训练功能.具体如下
class Featuremodel(tf.keras.Model): # 提取图片特征
def __init__(self ,intermediate_dim,latent_dim, **kwargs):
super(Featuremodel, self).__init__(**kwargs)
self.hidden_layer = Dense(units=intermediate_dim, activation=tf.nn.relu)
def call(self, x):
activation = self.hidden_layer(x)
return x,activation
class Autoencoder(tf.keras.Model):
def __init__(self, intermediate_dim, original_dim,latent_dim):
super(Autoencoder, self).__init__()
self.featuremodel = Featuremodel(intermediate_dim,latent_dim)
self.encoder = Encoder(intermediate_dim,latent_dim)
self.decoder = Decoder(intermediate_dim,original_dim)
def call(self, input_features):
x,feature = self.featuremodel(input_features)
z_mean,z_log_var = self.encoder(feature)
epsilon = K.random_normal(shape=(K.shape(z_mean)[0], latent_dim), mean=0.,
stddev=1.0)
code = z_mean + K.exp(z_log_var / 2) * epsilon
reconstructed = self.decoder(code)
return reconstructed,z_mean,z_log_varFeaturemodel类可以被扩展成更复杂的模型. Autoencoder则专注于变分训练.
6 配套资源下载方式
本文只是对tf2的基本使用做了简单的总结.全面系统的教程还要以书为参.另外tf2在BN的支持上也存在许多不便之处,例如，使用动态图训练时，可以为每个BN加入一个istraining参数，来控制模型是否需要更新BN中的均值和方差（因为在测试时不需要更新）；如果在keras模型体系中，则通过设置模型的trainable来控制。篇幅有限,这里不再展开.读者可以自行查找相关资料,在使用时还需多加小心.
腾讯云开发者
