test_general.py 文件源码

def test_export_model(self):
        input = Input(shape=(3, 4, 4), name='input', dtype='float32')
        out = Convolution2D(3, 3, 3, 
            init='he_normal', subsample=(1, 1), border_mode='same', activation='linear')(input)
        out = Activation('relu')(out)
        model = Model(input=[input], output=[out])

        data_model_folder = dir + "/../fixture/model_export"
        if K._BACKEND == 'tensorflow':
            import tensorflow as tf
            saver = tf.train.Saver()
        else:
            saver = None
        export_model(model, data_model_folder, saver=saver)

        os.remove(data_model_folder + '/archi.json')
        os.remove(data_model_folder + '/last_weights.hdf5')
        if K._BACKEND == 'tensorflow':
            os.remove(data_model_folder + '/checkpoint')
            os.remove(data_model_folder + '/tf-last_weights')
            os.remove(data_model_folder + '/tf-last_weights.meta')
            os.remove(data_model_folder + '/tf-model_graph')
            os.remove(data_model_folder + '/tf-frozen_model.pb')
        os.rmdir(data_model_folder)