Visualisations of model with letsplot (#70)

* introduce proposed plot api to jvm example

* fix docs

* rename function and refactor common code

* update plot api proposition

* create separate modeule for visualization

* naming fixes

* use it instead of model

* rename helper function xyio

* add kdocs

* add wav visualization

* add network visualization

* Add comments

* exclude slf4j, mark casts

* clean gradle configuration

* fix dependency in tests

Author: avan1235

api/src/main/kotlin/org/jetbrains/kotlinx/dl/dataset/embeddedDatasets.kt

@@ -21,12 +21,6 @@ import java.nio.file.Paths
 import java.nio.file.StandardCopyOption
 import java.util.zip.ZipEntry
 import java.util.zip.ZipFile
-import java.io.File
-
-import java.io.IOException
-
-import java.io.FileOutputStream
-import java.lang.IllegalStateException
 
 
 /**
@@ -191,7 +185,7 @@ public fun freeSpokenDigits(
     cacheDirectory.existsOrMkdirs()
 
     val path = freeSpokenDigitDatasetPath(cacheDirectory)
-    val dataset = File("$path/free-spoken-digit-dataset-master/recordings")
+    val dataset = File(path)
         .listFiles()?.flatMap(::extractWavFileSamples)
         ?: throw IllegalStateException("Cannot find Free Spoken Digits Dataset files in $path")
     val maxDataSize = dataset.map { it.first.size }.maxOrNull()
@@ -275,7 +269,8 @@ public fun dogsCatsDatasetPath(cacheDirectory: File = File("cache")): String =
     unzipDatasetPath(
         cacheDirectory,
         loadFile(cacheDirectory, DOGS_CATS_IMAGES_ARCHIVE),
-        "/datasets/dogs-vs-cats")
+        "/datasets/dogs-vs-cats"
+    )
 
 /** Path to the subset of Dogs-vs-Cats dataset. */
 private const val DOGS_CATS_SMALL_IMAGES_ARCHIVE: String = "datasets/small_catdogs/data.zip"
@@ -285,20 +280,25 @@ public fun dogsCatsSmallDatasetPath(cacheDirectory: File = File("cache")): Strin
     unzipDatasetPath(
         cacheDirectory,
         loadFile(cacheDirectory, DOGS_CATS_SMALL_IMAGES_ARCHIVE),
-        "/datasets/small-dogs-vs-cats")
+        "/datasets/small-dogs-vs-cats"
+    )
 
 /** Path to the Free Spoken Digits Dataset. */
 private const val FSDD_SOUNDS_ARCHIVE: String = "datasets/fsdd.zip"
 
 /** Path to download the Free Spoken Digits Dataset. */
-private const val FSS_SOUNDS_SOURCE: String = "https://codeload.github.com/Jakobovski/free-spoken-digit-dataset/zip/refs/heads/master"
+private const val FSS_SOUNDS_SOURCE: String =
+    "https://codeload.github.com/Jakobovski/free-spoken-digit-dataset/zip/refs/heads/master"
 
-/** Returns path to images of the subset of the Dogs-vs-Cats dataset. */
+/** Returns path to sound data files from Free Spoken Digits Dataset. */
 public fun freeSpokenDigitDatasetPath(cacheDirectory: File = File("cache")): String =
     unzipDatasetPath(
         cacheDirectory,
         loadFile(cacheDirectory, FSDD_SOUNDS_ARCHIVE, downloadURLFromRelativePath = { FSS_SOUNDS_SOURCE }),
-        "/datasets/free-spoken-digit")
+        "/datasets/free-spoken-digit"
+    ).run {
+        "$this/free-spoken-digit-dataset-master/recordings"
+    }
 
 /**
  * Download the compressed dataset from external source, decompress the file and remove the downloaded file

api/src/main/kotlin/org/jetbrains/kotlinx/dl/dataset/sound/wav/WavFile.kt

@@ -28,6 +28,11 @@ private const val RIFF_TYPE_ID: Long = 0x45564157
  *
  * Based on code written by [Andrew Greensted](http://www.labbookpages.co.uk/)
  * but modified to more Kotlin idiomatic way with only read option for simplicity.
+ *
+ * @property bufferSize is a size of a buffer to read from given file when reading next frames.
+ * @constructor creates [WavFile]
+ *
+ * @param file to read the WAV file data from
  */
 public class WavFile(
     file: File,

build.gradle

@@ -28,7 +28,8 @@ repositories {
 subprojects {
     repositories {
         mavenCentral()
+        maven {
+            url "https://dl.bintray.com/kotlin/kotlinx.html"
+        }
     }
 }
-
-

examples/build.gradle

@@ -1,12 +1,14 @@
 dependencies {
     implementation "org.jetbrains.kotlin:kotlin-stdlib-jdk8"
-    api(project(":api"))
+    api project(":api")
+    implementation project(":visualization")
 }
 
 dependencies {
     compile 'org.apache.logging.log4j:log4j-api:2.14.0'
     compile 'org.apache.logging.log4j:log4j-core:2.14.0'
     compile 'org.apache.logging.log4j:log4j-slf4j-impl:2.14.0'
+
     testCompile 'org.junit.jupiter:junit-jupiter-api:5.5.2'
     testCompile 'org.junit.jupiter:junit-jupiter-engine:5.5.2'
     testCompile 'org.junit.jupiter:junit-jupiter-params:5.5.2'
@@ -19,6 +21,7 @@ dependencies {
 compileKotlin {
     kotlinOptions.jvmTarget = "1.8"
 }
+
 compileTestKotlin {
     kotlinOptions.jvmTarget = "1.8"
 }
@@ -29,4 +32,3 @@ test {
     minHeapSize = "1024m"
     maxHeapSize = "8g"
 }
-

examples/src/main/kotlin/examples/cnn/fsdd/SoundNet.kt

@@ -37,7 +37,7 @@ private const val SEED = 12L
  * @param poolStride stride for poolSize and stride in maxpooling layer
  * @return array of layers to be registered in [Sequential] as vararg
  */
-private fun soundBlock(filters: Long, kernelSize: Long, poolStride: Long): Array<Layer> =
+internal fun soundBlock(filters: Long, kernelSize: Long, poolStride: Long): Array<Layer> =
     arrayOf(
         Conv1D(
             filters = filters,

examples/src/main/kotlin/examples/visualisation/LeNetFashionMnistVisualisation.kt renamed to examples/src/main/kotlin/examples/visualization/LeNetFashionMnistVisualization.kt

@@ -3,13 +3,16 @@
  * Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE.txt file.
  */
 
-package examples.visualisation
+package examples.visualization
 
 import examples.inference.lenet5
+import org.jetbrains.kotlinx.dl.api.core.layer.convolutional.Conv2D
 import org.jetbrains.kotlinx.dl.api.core.loss.Losses
 import org.jetbrains.kotlinx.dl.api.core.metric.Metrics
 import org.jetbrains.kotlinx.dl.api.core.optimizer.Adam
 import org.jetbrains.kotlinx.dl.dataset.fashionMnist
+import org.jetbrains.kotlinx.dl.visualization.letsplot.*
+import org.jetbrains.kotlinx.dl.visualization.swing.*
 
 private const val EPOCHS = 1
 private const val TRAINING_BATCH_SIZE = 500
@@ -38,7 +41,12 @@ fun main() {
 
     val (newTrain, validation) = train.split(0.95)
 
+    val sampleIndex = 42
+    val x = test.getX(sampleIndex)
+    val y = test.getY(sampleIndex).toInt()
+
     lenet5().use {
+
         it.compile(
             optimizer = Adam(),
             loss = Losses.SOFT_MAX_CROSS_ENTROPY_WITH_LOGITS,
@@ -53,28 +61,39 @@ fun main() {
             validationBatchSize = TEST_BATCH_SIZE
         )
 
-        val imageId = 0
-
-        val weights = it.layers[1].weights // first conv2d layer
-
-        drawFilters(weights.values.toTypedArray()[0], colorCoefficient = 10.0)
-
-        val weights2 = it.layers[3].weights // first conv2d layer
-
-        drawFilters(weights2.values.toTypedArray()[0], colorCoefficient = 12.0)
+        val fashionPlots = List(3) { imageIndex ->
+            flattenImagePlot(imageIndex, test,
+                predict = it::predict,
+                labelEncoding = fashionMnistLabelEncoding::get,
+                plotFeature = PlotFeature.GRAY
+            )
+        }
+        columnPlot(fashionPlots, 3, 256).show()
 
         val accuracy = it.evaluate(dataset = test, batchSize = TEST_BATCH_SIZE).metrics[Metrics.ACCURACY]
-
         println("Accuracy $accuracy")
 
-        val (prediction, activations) = it.predictAndGetActivations(train.getX(imageId))
+        val fstConv2D = it.layers[1] as Conv2D
+        val sndConv2D = it.layers[3] as Conv2D
 
-        println("Prediction: ${fashionMnistLabelEncoding[prediction]}")
+        // lets-plot approach
+        filtersPlot(fstConv2D, columns = 16).show()
+        filtersPlot(sndConv2D, columns = 16).show()
 
-        drawActivations(activations)
+        // swing approach
+        drawFilters(fstConv2D.weights.values.toTypedArray()[0], colorCoefficient = 10.0)
+        drawFilters(sndConv2D.weights.values.toTypedArray()[0], colorCoefficient = 10.0)
 
-        val trainImageLabel = train.getY(imageId)
+        val layersActivations = modelActivationOnLayersPlot(it, x)
+        val (prediction, activations) = it.predictAndGetActivations(x)
+        println("Prediction: ${fashionMnistLabelEncoding[prediction]}")
+        println("Ground Truth: ${fashionMnistLabelEncoding[y]}")
+
+        // lets-plot approach
+        layersActivations[0].show()
+        layersActivations[1].show()
 
-        println("Ground Truth: ${fashionMnistLabelEncoding[trainImageLabel.toInt()]}")
+        // swing approach
+        drawActivations(activations)
     }
 }

examples/src/main/kotlin/examples/visualisation/LeNetMnistVisualisation.kt renamed to examples/src/main/kotlin/examples/visualization/LeNetMnistVisualization.kt

@@ -3,13 +3,16 @@
  * Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE.txt file.
  */
 
-package examples.visualisation
+package examples.visualization
 
 import examples.inference.lenet5
+import org.jetbrains.kotlinx.dl.api.core.layer.convolutional.Conv2D
 import org.jetbrains.kotlinx.dl.api.core.loss.Losses
 import org.jetbrains.kotlinx.dl.api.core.metric.Metrics
 import org.jetbrains.kotlinx.dl.api.core.optimizer.Adam
 import org.jetbrains.kotlinx.dl.dataset.mnist
+import org.jetbrains.kotlinx.dl.visualization.letsplot.*
+import org.jetbrains.kotlinx.dl.visualization.swing.*
 
 private const val EPOCHS = 1
 private const val TRAINING_BATCH_SIZE = 500
@@ -21,11 +24,15 @@ private const val TEST_BATCH_SIZE = 1000
  * Model is trained on Mnist dataset.
  */
 fun main() {
+
     val (train, test) = mnist()
 
-    val imageId = 1
+    val sampleIndex = 42
+    val x = test.getX(sampleIndex)
+    val y = test.getY(sampleIndex).toInt()
 
     lenet5().use {
+
         it.compile(
             optimizer = Adam(),
             loss = Losses.SOFT_MAX_CROSS_ENTROPY_WITH_LOGITS,
@@ -42,26 +49,35 @@ fun main() {
             validationBatchSize = TEST_BATCH_SIZE
         )
 
-        val weights = it.layers[1].weights // first conv2d layer
-
-        drawFilters(weights.values.toTypedArray()[0], colorCoefficient = 10.0)
-
-        val weights2 = it.layers[3].weights // first conv2d layer
-
-        drawFilters(weights2.values.toTypedArray()[0], colorCoefficient = 10.0)
+        val numbersPlots = List(3) { imageIndex ->
+            flattenImagePlot(imageIndex, test, it::predict)
+        }
+        columnPlot(numbersPlots, 3, 256).show()
 
         val accuracy = it.evaluate(dataset = test, batchSize = TEST_BATCH_SIZE).metrics[Metrics.ACCURACY]
-
         println("Accuracy $accuracy")
 
-        val (prediction, activations) = it.predictAndGetActivations(train.getX(imageId))
+        val fstConv2D = it.layers[1] as Conv2D
+        val sndConv2D = it.layers[3] as Conv2D
 
-        println("Prediction: $prediction")
+        // lets-plot approach
+        filtersPlot(fstConv2D, columns = 16).show()
+        filtersPlot(sndConv2D, columns = 16).show()
 
-        drawActivations(activations)
+        // swing approach
+        drawFilters(fstConv2D.weights.values.toTypedArray()[0], colorCoefficient = 10.0)
+        drawFilters(sndConv2D.weights.values.toTypedArray()[0], colorCoefficient = 10.0)
 
-        val trainImageLabel = train.getY(imageId)
+        val layersActivations = modelActivationOnLayersPlot(it, x)
+        val (prediction, activations) = it.predictAndGetActivations(x)
+        println("Prediction: $prediction")
+        println("Ground Truth: $y")
 
-        println("Ground Truth: $trainImageLabel")
+        // lets-plot approach
+        layersActivations[0].show()
+        layersActivations[1].show()
+
+        // swing approach
+        drawActivations(activations)
     }
 }