[FIXED] prediction using custom generator returns different array size

Issue

I have a model that receives 3 inputs and gives 1 output (the model right now is just in order to have the same layers and sizes as in my real problem).

class CNNmodel():
    def __init__(self,
                 height,
                 width,
                 channels):
        self.height = height
        self.width = width
        self.channels = channels

    def build_model(self):
        inputs1 = Input((self.height, self.width, self.channels))
        inputs2 = Input((self.height, self.width, self.channels))
        inputs3 = Input((self.height, self.width, self.channels))
        
        M1M2 = Subtract()([inputs2, inputs1])
        M1M2E1 = Add()([M1M2, inputs3])
        concat = Concatenate(axis=0)([M1M2, inputs3])
        
        x = Conv2D(32, 1, activation='relu')(inputs1)
        
        sr1 = Conv2D(32, 1, activation='relu')(inputs2)
        sr2 = Conv2D(32, 1, activation='relu')(x)
        sr2 = Conv2D(32, 1, activation='relu')(sr2)
        sr2 = Conv2D(32, 1, activation='relu')(concat)

        addition = Add()([sr1, sr2])
       
        f1_input = Input((self.height, self.width, 32))
        f1 = Conv2D(32, 1, activation='relu')(f1_input)
        f2 = Conv2D(32, 1, activation='relu')(f1)
        f3 = Conv2D(1, 1,  activation='relu')(addition)
        
        
        outputs = Conv2D(3, 1)(f3)
        outputs = Conv2D(3, 1)(outputs)
        
        model = Model([inputs1, inputs2, inputs3], outputs)
        return model

I am using this generator:

class JoinedGen(tf.keras.utils.Sequence):
    def __init__(self,
                 input_gen1,
                 input_gen2,
                 input_gen3,
                 target_gen,
                 batch_size,
                 preds=False,
                 shuffle=False):
        self.gen1 = input_gen1
        self.gen2 = input_gen2
        self.gen3 = input_gen3
        self.target_gen = target_gen
        self.batch_size = batch_size
        self.preds = preds
        self.shuffle = shuffle
        self.on_epoch_end()

        assert len(input_gen1) == len(target_gen)

    def __len__(self):
        return int(np.floor(len(self.gen1) / self.batch_size))

    
    def __getitem__(self, i):
       
        x1 = self.gen1[i]
        x2 = self.gen2[i]
        x3 = self.gen3[i]
        y = self.target_gen[i]
       
        x1 = x1[np.newaxis, :, :, :]
        x2 = x2[np.newaxis, :, :, :]
        x3 = x3[np.newaxis, :, :, :]
        y = y[np.newaxis, :, :, :]
        
        if self.preds:
            return [x1, x2, x3]
        return [x1, x2, x3], y

    def on_epoch_end(self):
        self.indices = np.arange(len(self.gen1))
        if self.shuffle:
            np.random.shuffle(self.indices)

In order to predict, I am calling it with preds=True.

My data is:

import numpy as np
import tensorflow as tf
from tensorflow.keras.layers import Input, Conv2D, Add, Concatenate, Subtract
from tensorflow.keras.models import Model

x1 = np.random.random((64, 16, 16, 3))
x2 = np.random.random((64, 16, 16, 3))
x3 = np.random.random((64, 16, 16, 3))
y =  np.random.random((64, 16, 16, 3))

mygen = JoinedGen(x1, x2, x3, y, 16)

model = CNNmodel(16, 16, 3)
model = model.build_model()

loss = tf.keras.losses.Huber()
metric = ['accuracy']
optimizer = tf.keras.optimizers.Adam(learning_rate=0.0001)

model.compile(optimizer=optimizer,
              loss=loss,
              metrics=metric)


history = model.fit(mygen,
                    epochs=2,
                    batch_size=16)

mygen_preds = JoinedGen(x1, x2, x3,y, 16, preds=True)

Now, if I call :

preds = model.predict(mygen_preds)

I receive preds shape: (8, 16, 16, 3)

If I call:

preds = model.predict([mygen_preds.gen1, mygen_preds.gen2, mygen_preds.gen3],
                      batch_size=1)

I receive shape: (128, 16, 16, 3)

but the shape I expect is : (64, 16, 16, 3) since mygen_preds.gen1.shape (and all rest arrays) is (64, 16, 16, 3)

Solution

First, you are currently not specifying the batch size in your get_item function.

You are currently only adding a new dimension as batch dimension, but you probably should find a good way to specify the batch dimension.

This relates to the first dimension differing all the time while your batch size is constant.

def __getitem__(self, i):
    # Find batch start and end
    start = self.batch_size * i
    end = self.batch_size * (i + 1)

    # Ensure you dont out of range index the indices
    if end >= len(self.indices): end = len(self.indices) - 1

    x1 = self.gen1[start:end, :, :, :]
    x2 = self.gen2[start:end, :, :, :]
    x3 = self.gen3[start:end, :, :, :]
    y = self.target_gen[start:end, :, :, :]
    
    if self.preds:
        return [x1, x2, x3]
    return [x1, x2, x3], y

The expected output should be [16,16,16,3] as the first dimension should be equal to your batch dimension.

---

After comment edit;

The incompatable shapes comes from your concatenation on the batch dimension here:

    M1M2 = Subtract()([inputs2, inputs1])
    M1M2E1 = Add()([M1M2, inputs3])
    concat = Concatenate(axis=0)([M1M2, inputs3])

Change to this:

    concat = Concatenate(axis=-1)([M1M2, inputs3])

This way you concatenate over the channel dimension instead of the batch dimension and then can run the regular operations over an extended channel dimension.

Answered By - Albin Sidås