Try the following:
def build_model90():
input_layer = Input(shape = (2,))
dense_1 = Dense(units='1024', activation='relu')(input_layer)
dense_2 = Dense(units='512', activation='relu')(dense_1)
# NOTE: yields 4 outputs directly each with size 51
y_output = Dense(units=4 * 51, name='output')(dense_2)
y_output = Reshape((4, 51))(y_output)
# Define the model with the input layer and a list of output layers
model = tf.keras.Model(inputs=input_layer, outputs=y_output)
return model
forward_90 = build_model90()
optimizer = "Adam"
# NOTE: custom MSE loss
def mse_loss(x, y):
# x and y have shape (B, 4, 51)
error = tf.square(x - y)
# sum over 4 and 51 dimensions
loss = tf.reduce_sum(error, axis=[1, 2])
# average over batch size
return tf.reduce_mean(loss)
forward_90.compile(optimizer=optimizer, loss=mse_loss)
# NOTE: data reshaping
def reshape_data(x, y):
return tf.reshape(x, shape=(-1, 2)), \
tf.reshape(y, shape=(-1, 4, 51))
# prepare training and validation data
x_train, y_train = reshape_data(repeated_norm_train_dime, train_resp)
x_valid, y_valid = reshape_data(norm_test_dime, test_resp)
# Train the model for 100 epochs
epochs = 100
history_90 = forward_90.fit(x_train, y_train, epochs=epochs, batch_size=4040,
validation_data=(x_valid, y_valid))
Explaination:
- I redefined the model architecture to output one tensor (and not four) with shape
(4,51)
. After training you can recover individual tensors by indexing, like y_i = output[:, i]
.
- Then I defined a custom MSE loss, that sums over the dimensions
4
and 51
averaging over the batch size.
- Then, to be sure of the shape of the data I reshape them to
(N, 2)
and (N,4,51)
.
Let me know if this works for you.