when I change data set I receive: Layer sequential weight shape (3, 3, 1, 64) is not compatible with provided weight shape (4, 4, 3, 64)

Question

I'm new to py and machine-learning.

for university project I'm implementing a simple and basic federated cnn that should work on mnist, cifar10 and cifar100. Aside from the fact that for the moment I'm fine with the fact that all clients have the exact same data because I haven't done the partitioning yet, when you run the flower cifar10 simulation everything is fine but as soon as I change the dataset I receive "Layer sequential weight shape (3, 3, 1, 64) is not compatible with provided weight shape (4, 4, 3, 64)" I look at mnist and get "Layer sequential weight shape (256, 100) is not compatible with provided weight shape (256, 10)." about cifar10. what am I doing wrong? the model seems to me to be correctly constructed with variable values that change based on the dataset.

This is my code, or if you prefer https://github.com/FrancolinoFlavioDomenico/firstLevelDegreeThesis

main.py

import keras

import Server
import Model
import Plotter

from keras.datasets import mnist
from keras.datasets import cifar10
from keras.datasets import cifar100

from flwr.simulation.ray_transport.utils import enable_tf_gpu_growth

import time

from functools import lru_cache


enable_tf_gpu_growth()

# cifar10
# Model.Model.set_data(cifar10, 10, (4, 4), (32, 32, 3))
# Plotter.Plotter.dataset_name = 'Cifar10'
# server_cifar10 = Server.Server()
# server_cifar10.start_simulation()
# del server_cifar10
# print("finish cifar10")
# time.sleep(10)
# print("starting next dataset")


# Mnist
Model.Model.set_data(mnist, 10, (3, 3), (28, 28, 1))
Plotter.Plotter.dataset_name = 'mnist'
server_mnist = Server.Server()
server_mnist.start_simulation()
del server_mnist
print("finish mnist")
time.sleep(10)
print("starting next dataset")


# cifar100
Model.Model.set_data(cifar100, 100, (4, 4), (32, 32, 3))
Plotter.Plotter.dataset_name = 'Cifar100'
server_cifar100 = Server.Server()
server_cifar100.start_simulation()
del server_cifar100

model.py

import keras.src.backend
from keras.models import Sequential
from keras.layers import Conv2D, Dense, Flatten, Dropout
from keras.layers import MaxPooling2D

from keras.utils import to_categorical

from keras.datasets import cifar10

class Model:
    CLIENTS_NUM = 5
    ROUNDS_NUM = 3
    CLASSES_NUMBER = 10
    DATASET = cifar10
    KERNEL_SIZE = (4, 4)
    INPUT_SHAPE = (32, 32, 3)

    X_TRAIN = None
    Y_TRAIN = None
    X_TEST = None
    Y_TEST = None

    @classmethod
    def set_data(cls, dataset=None, classes_number=None, kernel_size=None, input_shape=None):
        if dataset is not None:
            Model.DATASET = dataset
            Model.CLASSES_NUMBER = classes_number
            Model.KERNEL_SIZE = kernel_size
            Model.INPUT_SHAPE = input_shape
            Model.DATASET = dataset
        # load data
        (Model.X_TRAIN, Model.Y_TRAIN), (Model.X_TEST, Model.Y_TEST) = Model.DATASET.load_data()
        # normalizing e label encodig
        Model.X_TRAIN = Model.X_TRAIN / 255
        Model.X_TEST = Model.X_TEST / 255
        Model.Y_TRAIN = to_categorical(Model.Y_TRAIN, Model.CLASSES_NUMBER)
        Model.Y_TEST = to_categorical(Model.Y_TEST, Model.CLASSES_NUMBER)

    @classmethod
    def get_model(cls):


        # build the model
        model = Sequential()
        model.add(Conv2D(64, Model.KERNEL_SIZE, input_shape=Model.INPUT_SHAPE, activation="relu"))
        model.add(MaxPooling2D(pool_size=(2, 2)))
        model.add(Dropout(0.5))
        model.add(Conv2D(64, Model.KERNEL_SIZE, input_shape=Model.INPUT_SHAPE, activation="relu"))
        model.add(MaxPooling2D(pool_size=(2, 2)))
        model.add(Dropout(0.25))
        model.add(Flatten())
        model.add(Dense(256, activation="relu"))
        model.add(Dense(Model.CLASSES_NUMBER, activation="softmax"))
        model.compile(loss="categorical_crossentropy", optimizer="adam", metrics=["accuracy"])
        return model

server.py

from typing import Dict, Optional, Tuple

import flwr as fl
import Model
from FlowerClient import get_client_fn

from flwr.simulation.ray_transport.utils import enable_tf_gpu_growth

import Plotter


class Server:

    def __init__(self) -> None:
        self.strategy = fl.server.strategy.FedAvg(
            min_fit_clients=Model.Model.CLIENTS_NUM,
            min_evaluate_clients=Model.Model.CLIENTS_NUM,
            min_available_clients=Model.Model.CLIENTS_NUM,
            evaluate_fn=self.get_eval_fn()
        )
        self.plotter = Plotter.Plotter()
        

    def start_simulation(self):
        client_resources = {"num_cpus": 2, "num_gpus": 0.25}
        fl.simulation.start_simulation(
            client_fn=get_client_fn(Model.Model.X_TRAIN, Model.Model.Y_TRAIN, Model.Model.X_TEST, Model.Model.Y_TEST),
            num_clients=Model.Model.CLIENTS_NUM,
            config=fl.server.ServerConfig(num_rounds=Model.Model.ROUNDS_NUM),
            strategy=self.strategy,
            client_resources=client_resources,
            actor_kwargs={
                "on_actor_init_fn": enable_tf_gpu_growth,
            }
        )

    def get_eval_fn(self):
        """Return an evaluation function for server-side evaluation."""
        # x_train, y_train, x_test, y_test = self.dataset

        """ x_train = self.xTrain
        y_train = self.yTrain """
        """ x_test = self.x_test
        y_test = self.y_test """

        # (x_train, y_train), (x_test, y_test) = ModelClass.getData()

        # The `evaluate` function will be called after every round
        def evaluate(
                server_round: int,
                parameters: fl.common.NDArrays,
                config: Dict[str, fl.common.Scalar],
        ) -> Optional[Tuple[float, Dict[str, fl.common.Scalar]]]:
            model = Model.Model.get_model()
            model.set_weights(parameters)  # Update model with the latest parameters
            loss, accuracy = model.evaluate(Model.Model.X_TEST, Model.Model.Y_TEST)
            self.plotter.accuracy_data.append(accuracy)
            self.plotter.loss_data.append(loss)
            if server_round == Model.Model.ROUNDS_NUM:
                self.plotter.plot()
            # print(f"After round {server_round}, Global accuracy = {accuracy}")
            # results = {"round":server_round,"loss": loss, "accuracy": accuracy}
            # results_list.append(results)
            return loss, {"accuracy": accuracy}

        return evaluate

client.py

import flwr as fl
import Model

from flwr.common import Metrics
from  keras.callbacks import EarlyStopping
from typing import Dict, List, Tuple


class FlowerClient(fl.client.NumPyClient):
    EPOCHS = 10
    BATCH_SIZE = 20
    STEPS_FOR_EPOCHS = 3
    VERBOSE = 0

    def __init__(self, x_train, y_train, x_test, y_test) -> None:
        self.x_train = x_train
        self.y_train = y_train
        self.x_test = x_test
        self.y_test = y_test
        self.model = Model.Model.get_model()

    def get_parameters(self, config):
        return self.model.get_weights()

    def fit(self, parameters, config):
        self.model.set_weights(parameters)
        self.model.fit(self.x_train, self.y_train, epochs=self.EPOCHS, batch_size=self.BATCH_SIZE,
                       steps_per_epoch=self.STEPS_FOR_EPOCHS)
        return self.model.get_weights(), len(self.x_train), {}

    def evaluate(self, parameters, config):
        self.model.set_weights(parameters)
        loss, accuracy = self.model.evaluate(self.x_test, self.y_test, verbose=self.VERBOSE)
        return loss, len(self.x_test), {"accuracy": float(accuracy)}


def get_client_fn(x_train, y_train, x_test, y_test):
    """Return a function to construct a client.

    The VirtualClientEngine will execute this function whenever a client is sampled by
    the strategy to participate.
    """

    def client_fn(cid: str) -> fl.client.Client:
        """Construct a FlowerClient with its own dataset partition."""
        # Extract partition for client with id = cid
        # client_dataset = dataset.load_partition(int(cid), "train")

        # # Now let's split it into train (90%) and validation (10%)
        # client_dataset_splits = client_dataset.train_test_split(test_size=0.1)

        # trainset = client_dataset_splits["train"].to_tf_dataset(
        #     columns="image", label_cols="label", batch_size = Cifar10Client.BATCH_SIZE
        # )
        # valset = client_dataset_splits["test"].to_tf_dataset(
        #     columns="image", label_cols="label", batch_size = Cifar10Client.BATCH_SIZE
        # )

        # trainset, valset = ModelClass.getData()

        # Create and return client

        # print(Client(m.Model.X_TRAIN, m.Model.Y_TRAIN, m.Model.X_TEST, m.Model.Y_TEST))
        return FlowerClient(x_train, y_train, x_test, y_test)

    return client_fn

As you can see from the comments in the main code, I also tried starting mnist directly without first going through cifar10 but to no avail, it's as if python/flowe/keras had a sort of cache in which once the data of the first dataset has been saved it always uses and only those