[Bug] Suspiciously slow calls to `Sequential.predictSoftly`

Calls to Sequential.predictSoftly start out fairly slow and each call takes a bit more time than the last one.

My example model is quite small, but the first call already takes about 25ms on my machine. (for comparison: fitting 1.000 random data points with 10 epochs takes 415ms) After calling predictSoftly 10.000 times, each successive call takes a third of a second.

Code to reproduce:

@OptIn(ExperimentalTime::class)
fun main() {
    Sequential.of(
        Input(36),
        Dense(36),
        Dense(36),
        Dense(36),
        Dense(36),
        Dense(36),
        Dense(16),
        Dense(8),
        Dense(3),
    ).use { model ->
        model.compile(
            optimizer = Adam(),
            loss = Losses.SOFT_MAX_CROSS_ENTROPY_WITH_LOGITS,
            metric = Metrics.MSE,
        )
        model.init()
        val features = FloatArray(36) { Random.nextFloat() }
        var predictionCalls = 0
        var predictionTimeOfBatch = Duration.ZERO
        val predictionTimes = mutableListOf<Double>()
        repeat(100_000) {
            val timing = measureTimedValue { model.predict(features) }
            predictionCalls++
            predictionTimeOfBatch += timing.duration
            predictionTimes += timing.duration.inMilliseconds
            if (predictionCalls % 100 == 0) {
                val csv = predictionTimes
                    .withIndex()
                    .joinToString("\n") { (i,t)->"${i + 1},$t" }
                File("timing.csv").writeText(csv)
                println("$predictionCalls calls done. (${predictionTimeOfBatch / 100} per call)")
                predictionTimeOfBatch = Duration.ZERO
            }
        }
    }
}

[zaleslaw]

Great idea for pet-project, so, hope it will be helpful, please create a separate ticket for this case as an feature request. So, I will release it in 0.1.1 but it takes time and will be released at the mid of January (not earlier, sorry)

[LostMekka]

c.) There should be a means to call predictSoftly on a batch of inputs. At the moment this is only possible for predict.

That would be awesome as well. My current pet project to try out this library is a board game AI that very crudely learns through self play, where the neural net is the search heuristic. The AI plays a semi-random game sequence and for each pair of successive moves I create a training data point. (needs 1 predictSoftly per data point) Building these self-play datasets would probably greatly benefit from batch soft-predicting 😄

Should I add a second issue for this?