model_util_test.py 文件源码

def test_model_pipe_keras(self):
        model = Sequential()
        model.add(Flatten(input_shape=(1, 28, 28)))
        model.add(Dense(128, activation='relu'))
        model.add(Dense(10, activation='softmax'))

        p = model_util.ModelPipe()
        input_data = [np.random.random((1, 1, 28, 28)) for _ in range(2)]

        p.add(model.predict, batch_size=64, batcher=np.vstack)

        expected_output = [
            model.predict(
                x.reshape(
                    (1, 1, 28, 28))) for x in input_data]
        output = p.apply_ordered(input_data)

        self.assertTrue(np.isclose(np.array(output).flatten(),
                                   np.array(expected_output).flatten()).all())