TrainingDivergedException executing TransferFreshFruit Example + Differences in ATLearn Embeddings

[rd-peter-braun]

I got
Caused by: ai.djl.TrainingDivergedException: The Loss became NaN, try reduce learning rate,add clipGradient option to your optimizer, check input data and loss calculation. at ai.djl.training.listener.DivergenceCheckTrainingListener.onTrainingBatch(DivergenceCheckTrainingListener.java:27) ~[api-0.28.0.jar:na]
Exception during the SoftmaxCrossEntropy evaluation when executing the TransferFreshFruit Example.
It occurs during the second epoch.

I tried it with the djl embedding model djl://ai.djl.pytorch/resnet18_embedding and with a self generated one using ATLearn.
Latter had a differing final layer - .addSingleton(nd -> nd.squeeze(new int[] {2, 3})) did not work since only 2 dimensions existed.

I use a custom trainset containing 83 classes - each class contains 250 images (jpg size 640x480). Which means my dataset construction differs from the DJL Example...

properties + code example:

trained-model-name: "eb_resnet_18" // props.getTrainedModelName()...
trainset-path: "trainset" // contains 83 classes....
epochs: 25
engine: PyTorch
device: CPU
batch-size: 32
learning-rate: 0.001
patience: 4
train-param: true

    public TrainingResult trainEmbeddingModel()
            throws IOException, TranslateException, ModelException, URISyntaxException {
        // Count the number of classes in the dataset
        long numClasses = Files.list(Paths.get(props.getTrainsetPath())).filter(Files::isDirectory).count();
        String modelUrl = "djl://ai.djl.pytorch/resnet18_embedding";
        Criteria<NDList, NDList> criteria =
                Criteria.builder()
                        .setTypes(NDList.class, NDList.class)
                        .optModelUrls(modelUrl)
                        .optEngine(props.getEngine())
                        .optProgress(new ProgressBar())
                        .optOption("trainParam", props.getTrainParam())
                        .build();

        ZooModel<NDList, NDList> embedding = criteria.loadModel();

        Block baseBlock = embedding.getBlock();
        Block blocks =
                new SequentialBlock()
                        .add(baseBlock)
                        .addSingleton(nd -> nd.squeeze(new int[] {2, 3}))
                        .add(Linear.builder().setUnits(numClasses).build())
                        .addSingleton(nd -> nd.softmax(1));

        Model model = Model.newInstance(props.getTrainedModelName(), props.getEngine());
        model.setBlock(blocks);

        // Configure trainer
        DefaultTrainingConfig config = setupTrainingConfig(baseBlock);
        Trainer trainer = model.newTrainer(config);
        trainer.setMetrics(new Metrics());

        // Initialize the parameter shape and value
        Shape inputShape = new Shape(props.getBatchSize(), 3, 224, 224);
        trainer.initialize(inputShape);

        // Data
        RandomAccessDataset[] splitDatasets = getRandomAccessDatasets();

        // Train
        EasyTrain.fit(trainer, props.getEpochs(), splitDatasets[0], splitDatasets[1]);

        // Save model
        model.save(Paths.get("SAVE_PATH"), props.getTrainedModelName());

        model.close();
        embedding.close();
        return trainer.getTrainingResult();
    }


    private DefaultTrainingConfig setupTrainingConfig(Block baseBlock) {
        SaveModelTrainingListener listener = new SaveModelTrainingListener(DATA_PATH, props.getTrainedModelName(), 1);
        listener.setSaveModelCallback(
                trainer -> {
                    TrainingResult result = trainer.getTrainingResult();
                    Model model = trainer.getModel();
                    float accuracy = result.getValidateEvaluation("Accuracy");
                    model.setProperty("Accuracy", String.format("%.5f", accuracy));
                    model.setProperty("Loss", String.format("%.5f", result.getValidateLoss()));
                });

        DefaultTrainingConfig config =
                new DefaultTrainingConfig(new SoftmaxCrossEntropy("SoftmaxCrossEntropy"))
                        .addEvaluator(new Accuracy())
//                        .addEvaluator(new TopKAccuracy(4))
                        .optDevices(Engine.getEngine(props.getEngine()).getDevices(1))
                        .addTrainingListeners(TrainingListener.Defaults.logging(DATA_PATH))
                        .addTrainingListeners(listener);

        // Customized learning rate
        float lr = props.getLearningRate();
        FixedPerVarTracker.Builder learningRateTrackerBuilder =
                FixedPerVarTracker.builder().setDefaultValue(lr);
        for (Pair<String, Parameter> paramPair : baseBlock.getParameters()) {
            learningRateTrackerBuilder.put(paramPair.getValue().getId(), 0.1f * lr);
        }
        FixedPerVarTracker learningRateTracker = learningRateTrackerBuilder.build();
        Optimizer optimizer = Adam.builder().optLearningRateTracker(learningRateTracker).build();
        config.optOptimizer(optimizer);

        return config;
    }
    private static class SoftmaxCrossEntropy extends Loss {
        /**
         * Base class for metric with abstract update methods.
         *
         * @param name The display name of the Loss
         */
        public SoftmaxCrossEntropy(String name) {
            super(name);
        }

        /** {@inheritDoc} */
        @Override
        public NDArray evaluate(NDList labels, NDList predictions) {
            // Here the labels are supposed to be one-hot
            int classAxis = -1;
            NDArray pred = predictions.singletonOrThrow().log();
            NDArray lab = labels.singletonOrThrow().reshape(pred.getShape());
            NDArray loss = pred.mul(lab).neg().sum(new int[] {classAxis}, true);
            return loss.mean();
        }
    }
    private RandomAccessDataset[] getRandomAccessDatasets() throws IOException, TranslateException {
        // Load dataset with transformations
        Pipeline pipeline = new Pipeline();
        pipeline.add(new Resize(224, 224))  // Resize images to 224x224
                .add(new ToTensor())
                .add(new Normalize(new float[]{0.485f, 0.456f, 0.406f}, new float[]{0.229f, 0.224f, 0.225f})); // Normalize to the same mean and std as used in ResNet

        // Load dataset
        ImageFolder dataset = ImageFolder.builder()
                .setRepositoryPath(Paths.get(props.getTrainsetPath()))
                .optPipeline(pipeline)  // Apply the transformations
                .addTargetTransform(new OneHot(83)) //FIXME: hardcoded number of classes
                .setSampling(props.getBatchSize(), true) // originally batch size was 600 with learning rate 0.007
                .build();
        dataset.prepare();

        // Split the dataset into training and validation sets
        RandomAccessDataset[] splitDatasets = dataset.randomSplit(9, 1, 0); // 90% training, 10% validation, 0% test
        return splitDatasets;
    }