MLP model training is slow

[konsteva]

Greetings!

I'm training a simple MLP model using DJL v0.18.0 on CPU. The model is trained successfully but the training seems to be rather slow taking around 7 minutes. My model is the following:

SequentialBlock block = new SequentialBlock();
block.add(Linear.builder()
    .setUnits(360)
    .build());
block.add(Activation::relu);
block.add(Linear.builder()
    .setUnits(128)
    .build());
block.add(Activation::relu);
block.add(Linear.builder()
    .setUnits(64)
    .build());
block.add(Activation::relu);
block.add(Linear.builder()
    .setUnits(32)
    .build());
block.add(Activation::relu);
block.add(Linear.builder()
    .setUnits(5)
    .build());

block.setInitializer(new XavierInitializer(), Type.WEIGHT);

Model mlpModel = Model.newInstance("MLP Model");
mlpModel.setBlock(block);

and the code for the training is the following:

DefaultTrainingConfig trainingConfig = new DefaultTrainingConfig(Loss.softmaxCrossEntropyLoss())
    .optOptimizer(Optimizer.adam().optEpsilon((float)0.001).build());

Trainer trainer = mlpModel.newTrainer(trainingConfig);
trainer.initialize(new Shape(360));

EasyTrain.fit(trainer, 500, myCustomTrainSet, null);

Some additional info about the training data:

• Custom ArrayDataset
• DataType.Float32
• 3850 instances in total
• Batch size 64

Note that I'm only counting the EasyTrain.fit method and nothing else. Are there any anti-patterns in my code that is making it slow? Is there anything else I can do to speed it up without changing the model architecture and hyperparameters?

Thanks in advance!