Wrong predict results using Tensorflow Java

Question

I am doing some experiments using TensorFlow Java and Python. I trained and saved a MNIST model using Python and now I want to load it in Java and make prediction. I have the following method

public void predict() {
    final int IMAGE_INDEX = 10;
    MnistDataset dataset = MnistDataset.create(0, TRAINING_IMAGES_ARCHIVE, TRAINING_LABELS_ARCHIVE,
            TEST_IMAGES_ARCHIVE, TEST_LABELS_ARCHIVE);

    try (SavedModelBundle modelBundle = SavedModelBundle
            .load("PATH_TO_TRAINED_MODEL/mymnist");
         Tensor batchImages = TUint8.tensorOf(dataset.testBatch().images().get(IMAGE_INDEX))) {
        FloatDataBuffer floatDataBuffer = DataBuffers.ofFloats(784);
        float[] ds = new float[784];
        for (int i = 0; i < batchImages.rawData().size(); i++) {
            ds[i] = batchImages.rawData().getByte(i);
        }
        floatDataBuffer = floatDataBuffer.write(ds);

        try (Tensor input = TFloat32.tensorOf(Shape.of(1, 28, 28, 1), floatDataBuffer);
             Tensor tensor = modelBundle.session().runner()
                     .feed("serving_default_input", input)
                     .fetch("StatefulPartitionedCall")
                     .run().get(0).expect(TFloat32.DTYPE)) {
            System.out.print(argmax(tensor.data().get(0)));
            System.out.println("-" + dataset.testBatch().labels().get(IMAGE_INDEX).getByte());

        }
    }
}

The model in Python looks like the following

model = Sequential()
model.add(Input(shape=(28, 28, 1), dtype="float32", name="input"))
model.add(Conv2D(filters=32, kernel_size=(4, 4), input_shape=(28, 28, 1), activation='relu'))
model.add(MaxPool2D(pool_size=(2, 2)))
model.add(Flatten())
model.add(Dense(256, activation='relu'))
model.add(Dense(10, activation='softmax'))
model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])
early_stopping = EarlyStopping(monitor='val_loss', patience=2)
history = model.fit(x_train, y_train_categorical, epochs=10, validation_data=(x_test, y_test_categorical),
                    callbacks=[early_stopping])

The input layer is float32 but MNIST images are Tensor so I had to convert them to Tensor to be able to feed the model in Java. I am not sure if this conversion is a good idea which I came up with (it bugs me ;))

    FloatDataBuffer floatDataBuffer = DataBuffers.ofFloats(784);
    float[] ds = new float[784];
    for (int i = 0; i < batchImages.rawData().size(); i++) {
        ds[i] = batchImages.rawData().getByte(i);
    }
    floatDataBuffer = floatDataBuffer.write(ds);

    Tensor input = TFloat32.tensorOf(Shape.of(1, 28, 28, 1), floatDataBuffer);

I get always incorrect result in Java code but when I call predict on the original model in Python, it works fine. Could someone please help me to resolve the problem?

Thanks in advance

PS. I used MnistDataset

Wrong predict results using Tensorflow Java

0 Answers