LpipsSmokeTest.kt

    private fun texturedPairs(size: Int): List<Pair> {
        val n = size * size
        fun field(f: (x: Int, y: Int) -> Float) = FloatArray(n) { f(it % size, it / size) }
        fun rgb(r: FloatArray, g: FloatArray, b: FloatArray): Img {
            val data = FloatArray(n * 3)
            for (p in 0 until n) { data[p * 3] = r[p]; data[p * 3 + 1] = g[p]; data[p * 3 + 2] = b[p] }
            return Img(size, size, 3, data)
        }
        val base = pairs(size)
        val noise = base[3].a
        val smooth = rgb(
            field { x, y -> 0.5f + 0.4f * kotlin.math.sin(0.37f * x + 0.21f * y) },
            field { x, y -> 0.5f + 0.35f * kotlin.math.cos(0.19f * x - 0.43f * y) },
            field { x, y -> 0.45f + 0.3f * kotlin.math.sin(0.11f * x * y / size + 0.3f * x) }
        )
        return listOf(
            Pair("noise-vs-disc@$size", noise, base[0].a),
            Pair("noise-vs-checker@$size", noise, base[2].b),
            Pair("smooth-vs-disc@$size", smooth, base[0].a),
            Pair("smooth-vs-ramp@$size", smooth, base[1].b)
        )
    }

    private fun gradient() {
        val started = System.nanoTime()
        Lpips.ensureGradModel()
        println("[lpips] gradient model ready in %.1f s".format((System.nanoTime() - started) / 1e9))
        val grad = Lpips.gradSession
        check(
            grad.inputNames == listOf("candidate", "reference") && grad.outputNames == listOf("distance", "gradient"),
            "gradient graph signature is (candidate, reference) -> (distance, gradient): ${grad.inputNames} -> ${grad.outputNames}"
        )
        // Its distance is the forward model's, on every fixed pair.
        var worstDistance = 0.0
        for (size in listOf(32, 64, 128)) for (p in pairs(size)) {
            val a = planarOf(p.a)
            val b = planarOf(p.b)
            val d = Lpips.distanceAndGradientPlanar(a, size, size, b).first
            worstDistance = maxOf(worstDistance, abs(d - Lpips.distancesPlanar(a, 1, size, size, b)[0]))
        }
        check(worstDistance <= 1e-6, "the gradient model's distance matches the forward model's on ${REFERENCE.size} pairs (max diff %.1e)".format(worstDistance))

        var worstPixel = 0.0
        var worstDirection = 0.0
        var values = 0
        var kinks = 0
        for (size in listOf(32, 64)) for (p in texturedPairs(size)) {
            val a = planarOf(p.a)
            val b = planarOf(p.b)
            val g = Lpips.distanceAndGradientPlanar(a, size, size, b).second
            fun at(x: FloatArray) = Lpips.distancesPlanar(x, 1, size, size, b)[0]
            // The sixteen values with the largest |gradient|, less any whose +-h step crosses a kink:
            // AlexNet is piecewise linear in a pixel, so on one linear piece the two one-sided
            // slopes agree, and where they do not, a central difference measures neither side.
            val h = FD_STEP.toFloat()
            val centre = at(a)
            var checked = 0
            // The smooth picture at 64 has a kink within +-1e-3 of most of its largest-gradient
            // values, some with one-sided slopes that happen to agree: in float64 PyTorch at (ch 1,
            // x 23, y 21) a step of 1e-3 gives 3.714e-3 where 1e-4 to 1e-6 and autograd give 3.798e-3
            // (and this gradient, 3.798e-3). It gets the directional check only.
            val singles = if (p.name.startsWith("smooth") && size == 64) emptyList() else g.indices.sortedByDescending { abs(g[it]) }.take(16)
            for (i in singles) {
                val plus = at(a.copyOf().also { it[i] += h })
                val minus = at(a.copyOf().also { it[i] -= h })
                val right = (plus - centre) / h
                val left = (centre - minus) / h
                if (abs(right - left) > FD_TOLERANCE * maxOf(abs(right), abs(left))) {
                    kinks++
                    continue
                }
                val fd = (plus - minus) / (2.0 * h)
                val rel = abs(fd - g[i]) / maxOf(abs(fd), abs(g[i].toDouble()))
                worstPixel = maxOf(worstPixel, rel)
                checked++
                check(rel <= FD_TOLERANCE, "${p.name} value $i (ch ${i / (size * size)}, x ${i % size}, y ${(i / size) % size}): gradient %.4e, central difference %.4e (rel %.1e)".format(g[i], fd, rel))
            }
            if (singles.isNotEmpty()) check(checked >= 4, "${p.name}: $checked of the 16 largest-gradient values lie on one linear piece within +-$h (need 4)")
            values += checked
            // Along the gradient's own direction: d(t) = d(a + t g / |g|), so d'(0) should be |g|.
            val norm = kotlin.math.sqrt(g.sumOf { it.toDouble() * it })
            val t = FD_STEP
            val plus = FloatArray(a.size) { a[it] + (t * g[it] / norm).toFloat() }
            val minus = FloatArray(a.size) { a[it] - (t * g[it] / norm).toFloat() }
            val fd = (at(plus) - at(minus)) / (2 * t)
            val rel = abs(fd - norm) / maxOf(abs(fd), norm)
            worstDirection = maxOf(worstDirection, rel)
            check(rel <= FD_TOLERANCE, "${p.name}: |gradient| %.4e against the central difference along it %.4e (rel %.1e)".format(norm, fd, rel))
        }
        println(
            "[lpips] gradient vs central difference (h = %.0e): worst relative %.2e on $values single values ($kinks skipped as straddling a kink), %.2e along the gradient on 8 pairs"
                .format(FD_STEP, worstPixel, worstDirection)
        )

        for (size in listOf(32, 64, 128)) {
            val p = pairs(size)[1]
            val a = planarOf(p.a)
            val b = planarOf(p.b)
            repeat(3) { Lpips.distanceAndGradientPlanar(a, size, size, b) }
            val t0 = System.nanoTime()
            repeat(10) { Lpips.distanceAndGradientPlanar(a, size, size, b) }
            println("[lpips] gradient cost @${size}x$size: %.2f ms for one pair".format((System.nanoTime() - t0) / 1e6 / 10))
        }
    }

    @JvmStatic
    fun main(args: Array<String>) {
        val started = System.nanoTime()
        val path = Lpips.ensureModel()
        println("[lpips] model $path (${java.nio.file.Files.size(path)} bytes), ready in %.1f s".format((System.nanoTime() - started) / 1e9))
        val session = Lpips.session
        check(
            session.inputNames == listOf("candidate", "reference") && session.outputNames == listOf("distance"),
            "graph signature is (candidate, reference) -> distance: ${session.inputNames} -> ${session.outputNames}"
        )

        var worst = 0.0
        for (size in listOf(32, 64, 128)) {
            val ps = pairs(size)
            for (p in ps) {
                val got = Lpips.distance(p.a, p.b)
                val want = REFERENCE.getValue(p.name)
                worst = maxOf(worst, abs(got - want))
                check(abs(got - want) <= TOLERANCE, "${p.name}: %.7f against PyTorch %.7f".format(got, want))
            }
            val byName = ps.associateBy { it.name }
            val same = Lpips.distance(byName.getValue("disc-vs-grey-disc@$size").a, byName.getValue("disc-vs-grey-disc@$size").b)
            val different = Lpips.distance(byName.getValue("disc-vs-ramp@$size").a, byName.getValue("disc-vs-ramp@$size").b)
            check(same < different, "@$size: a grey disc (%.3f) is nearer the yellow disc than a ramp is (%.3f)".format(same, different))
            check(Lpips.distance(ps[0].a, ps[0].a) == 0.0, "@$size: a picture is at distance 0 from itself")

            val reference = ps[0].a
            val candidates = ps.map { it.b }
            val batched = Lpips.distances(candidates, reference)
            val singles = candidates.map { Lpips.distance(it, reference) }
            val batchErr = batched.indices.maxOf { abs(batched[it] - singles[it]) }
            check(batchErr <= 1e-6, "@$size: ${candidates.size} candidates against one broadcast reference match one-at-a-time scoring (max diff %.1e)".format(batchErr))
        }
        println("[lpips] max |field - pytorch| = %.2e over ${REFERENCE.size} pairs".format(worst))

        // Cost, printed: per evaluation, alone and in a batch of 16 against one reference.
        for (size in listOf(32, 64, 128)) {
            val p = pairs(size)
            val batch = List(16) { p[it % p.size].b }
            repeat(3) { Lpips.distances(batch, p[0].a) }
            val t0 = System.nanoTime()
            repeat(20) { Lpips.distance(p[1].a, p[1].b) }
            val single = (System.nanoTime() - t0) / 1e6 / 20
            val t1 = System.nanoTime()
            repeat(5) { Lpips.distances(batch, p[0].a) }
            val perInBatch = (System.nanoTime() - t1) / 1e6 / (5 * 16)
            println("[lpips] cost @${size}x$size: %.2f ms for one pair, %.2f ms per candidate in a batch of 16".format(single, perInBatch))
        }

        gradient()

        if (failures.isNotEmpty()) {
            println("[lpips] ${failures.size} failure(s):")
            failures.forEach { println("  - $it") }
            exitProcess(1)
        }
        println("ONNX LPIPS smoke tests passed")
    }
}