tflearn重载预训练模型时的bug及修复
最近接触了tflearn框架,是一个集合了tensorflow的一些功能较为高层的API。模型搭建和训练也比tensorflow的语句精炼,但就是容易出bug后找不到原因。比如最近在重载预训练模型时就出现了一个莫名其妙的bug。google找不到类似的经历,进包自带的函数里调试才发现了问题所在,记录下来也许能帮助到有用的同志。
问题描述
最近在做的是从一个生物信号S里重建图像,思路就是先用AutoEncoder提图像特征ft,再把生物信号S往ft做回归,再从ft上采样,重建图像。用此种方法作为后续研究的baseline。
无奈手里与生物信号成对的图像数量很少,最多也只有不到2000,自己设计训练vae或者ae都效果不佳。于是往github上找模型,找到一个在MSCOCO+ Flickr 30k. 总共约 200k图像上进行预训练的autoencoder模型,地址。测试时可以接受jpg或者png图像,程序会reshape到256*256.train_encoder.py里的模型为:
def build_model():
# logging.info('building model')
img_prep = ImagePreprocessing()
img_prep.add_featurewise_zero_center()
img_prep.add_featurewise_stdnorm()
encoder = input_data(shape=(None, IMAGE_INPUT_SIZE[0], IMAGE_INPUT_SIZE[1],
3), data_preprocessing=img_prep)
encoder = conv_2d(encoder, 16, 7, activation='relu') # incoming, nb_filter, filter_size, strides=1, padding='same', activation='linear',
encoder = dropout(encoder, 0.25) # you can have noisy input instead
encoder = max_pool_2d(encoder, 2)
encoder = conv_2d(encoder, 16, 7, activation='relu')
encoder = max_pool_2d(encoder, 2)
encoder = conv_2d(encoder, 8, 7, activation='relu')
encoder = max_pool_2d(encoder, 2)
decoder = conv_2d(encoder, 8, 7, activation='relu')
decoder = upsample_2d(decoder, 2)
decoder = conv_2d(decoder, 16, 7, activation='relu')
decoder = upsample_2d(decoder, 2)
decoder = conv_2d(decoder, 16, 7, activation='relu')
decoder = upsample_2d(decoder, 2)
decoder = conv_2d(decoder, 3, 7)
return regression(decoder, optimizer='adadelta',
loss='binary_crossentropy', learning_rate=0.005)
evaluate.py里导入train_encoder.py里的build_model函数,并从对应的checkpoints里导入模型的权重参数
net = train_encoder.bulid_model()
model = tflearn.DNN(net)
checkpoint_path = './checkpoints/model.h5'
model.load(checkpoint_path)
我们也可以只重载encoder模块,这比后面的只重载decoder模块要容易:
encoder = input_data(shape=(None, IMAGE_INPUT_SIZE[0], IMAGE_INPUT_SIZE[1],
3))# , data_preprocessing=img_prep, name='image_input')
encoder = conv_2d(encoder, 16, 7, activation='relu') # incoming, nb_filter, filter_size, strides=1, padding='same', activation='linear',
encoder = dropout(encoder, 0.25) # you can have noisy input instead
encoder = max_pool_2d(encoder, 2)
encoder = conv_2d(encoder, 16, 7, activation='relu')
encoder = max_pool_2d(encoder, 2)
encoder = conv_2d(encoder, 8, 7, activation='relu')
encoder_model = tflearn.DNN(encoder, session=model.session)
在我们需要从自己的ft里,生成图像时,我们只需要重载decoder。只重载decoder模块时,问题就出现了。
z = input_data(shape=(None,32,32,8), name='z')
decoder = conv_2d(z, 8, 7, activation='relu', scope='Conv2D_3', reuse=True)
decoder = upsample_2d(decoder, 2, )
decoder = conv_2d(decoder, 16, 7, activation='relu', scope='Conv2D_4', reuse=True)
decoder = upsample_2d(decoder, 2, )
decoder = conv_2d(decoder, 16, 7, activation='relu', scope='Conv2D_5', reuse=True)
decoder = upsample_2d(decoder, 2, )
decoder = conv_2d(decoder, 3, 7, scope='Conv2D_6', reuse=True)
generator_model = tflearn.DNN(decoder, session=model.session)
recons_img = generator_model.predict({'z':ft})
ft就是我从生物信号S回归出的特征。作为生成模型的输入。但运行时出现了error
KeyError: <tf.Tensor 'InputData/X:0' shape=(?, 256, 256, 3) dtype=float32>
调试一下, 怎么generator_model的inputs里面有两个tensor?一个的shape是(None, 32,32,8)
,还有一个就是shape=(None,256,256,3)
是整个model的input_data层的shape。我以为是自己代码的编写有问题,但是参考了github上的tflearn教程里的一个例子,地址。发现里面的generator_model里也有两个tensor但是运行不会报错。郁闷了很久,最后还是跟着error Traceback进到产生错误的根源位置=>"/nfs/software/anaconda3/lib/python3.5/site-packages/tflearn/helpers/evaluator.py"
class Evaluator(object):
def __init__(self, tensors, model=None, session=None):
self.tensors = to_list(tensors)
self.graph = self.tensors[0].graph
self.model = model
self.dprep_collection = tf.get_collection(tf.GraphKeys.DATA_PREP)
self.inputs = tf.get_collection(tf.GraphKeys.INPUTS)
with self.graph.as_default():
self.session = tf.Session()
if session: self.session = session
self.saver = tf.train.Saver()
if model: self.saver.restore(self.session, model)
def predict(self, feed_dict):
with self.graph.as_default():
# Data Preprocessing
dprep_dict = dict()
for i in range(len(self.inputs)):
# Support for custom inputs not using dprep/daug
if len(self.dprep_collection) > i:
if self.dprep_collection[i] is not None:
dprep_dict[self.inputs[i]] = self.dprep_collection[i]
# Apply pre-processing
if len(dprep_dict) > 0:
for k in dprep_dict:
feed_dict[k] = dprep_dict[k].apply(feed_dict[k])
# Prediction for each tensor
tflearn.is_training(False, self.session)
prediction = []
if len(self.tensors) == 1:
return self.session.run(self.tensors[0], feed_dict=feed_dict)
else:
for output in self.tensors:
o_pred = self.session.run(output, feed_dict=feed_dict).tolist()
for i, val in enumerate(o_pred): # Reshape pred per sample
if len(self.tensors) > 1:
if not len(prediction) > i: prediction.append([])
prediction[i].append(val)
return prediction
注意这里我代码里的feed_dict
就是{'z':ft}
,明确指定了两个tensor
中的一个,但是为什么它就没有指定我指定的那个tensor
呢?发现feed_dict
在len(dprep_dict) > 0
时会有个改动,而这个dpre_dict
是否长度为0,则与self.dprep_collection[i] is not None
这个判断条件True还是False有关。调试时tflearn给的例子里,dprep_collection
都是None元素,而我的代码里有一个元素跟ImagePreprocessing
有关,仔细一看,model 的input_data里的一个参数是data_preprocessing=img_prep
,把它注释掉以后,error就消失了!
最后推荐一个入门tflearn的项目,地址放在这里。examples里面有很多简单的项目代码,可以从中学到如何搭建模型,如何训练模型,如何finetune等等。