Add PReLU layer

api/src/main/kotlin/org/jetbrains/kotlinx/dl/api/core/layer/activation/PReLU.kt

@@ -0,0 +1,83 @@
+/*
+ * Copyright 2020 JetBrains s.r.o. and Kotlin Deep Learning project contributors. All Rights Reserved.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlinx.dl.api.core.layer.activation
+
+import org.jetbrains.kotlinx.dl.api.core.KGraph
+import org.jetbrains.kotlinx.dl.api.core.initializer.Initializer
+import org.jetbrains.kotlinx.dl.api.core.initializer.Zeros
+import org.jetbrains.kotlinx.dl.api.core.regularizer.Regularizer
+import org.jetbrains.kotlinx.dl.api.core.shape.numElements
+import org.jetbrains.kotlinx.dl.api.core.shape.toLongArray
+import org.jetbrains.kotlinx.dl.api.core.util.getDType
+import org.tensorflow.Operand
+import org.tensorflow.Shape
+import org.tensorflow.op.Ops
+import org.tensorflow.op.core.Variable
+
+/**
+ * Parametric Rectified Linear Unit.
+ *
+ * It follows:
+ * ```
+ * f(x) = alpha * x     if x < 0
+ * f(x) = x             if x >= 0
+ * ```
+ * where `alpha` is a learnable weight and has the same shape as `x` (i.e. input).
+ *
+ * @property [alphaInitializer] Initializer instance for the weights.
+ * @property [alphaRegularizer] Regularizer instance for the weights.
+ * @property [sharedAxes] The axes along which to share learnable parameters.
+ */
+public class PReLU(
+    public val alphaInitializer: Initializer = Zeros(),
+    public val alphaRegularizer: Regularizer? = null,
+    public val sharedAxes: IntArray? = null,
+    name: String = ""
+) : AbstractActivationLayer(name) {
+    /**
+     * TODO: support for constraint (alphaConstraint) should be added
+     */
+
+    private lateinit var alphaShape: Shape
+    private lateinit var alpha: Variable<Float>
+    private val alphaVariableName = if (name.isNotEmpty()) name + "_" + "alpha" else "alpha"
+
+    override var weights: Map<String, Array<*>>
+        get() = extractWeights(listOf(alphaVariableName))
+        set(value) = assignWeights(value)
+    override val paramCount: Int
+        get() = alphaShape.numElements().toInt()
+
+    init {
+        isTrainable = true
+    }
+
+    override fun build(tf: Ops, kGraph: KGraph, inputShape: Shape) {
+        val alphaShapeArray = inputShape.toLongArray().drop(1).toLongArray()
+        if (sharedAxes != null) {
+            for (axis in sharedAxes) {
+                alphaShapeArray[axis - 1] = 1
+            }
+        }
+        alphaShape = Shape.make(alphaShapeArray[0], *alphaShapeArray.drop(1).toLongArray())
+
+        fanIn = inputShape.size(inputShape.numDimensions() - 1).toInt()
+        fanOut = fanIn
+
+        alpha = tf.withName(alphaVariableName).variable(alphaShape, getDType())
+        alpha = addWeight(tf, kGraph, alphaVariableName, alpha, alphaInitializer, alphaRegularizer)
+    }
+
+    override fun forward(tf: Ops, input: Operand<Float>): Operand<Float> {
+        // It's equivalent to: `-alpha * relu(-x) + relu(x)`
+        val positive = tf.nn.relu(input)
+        val negative = tf.math.mul(tf.math.neg(alpha), tf.nn.relu(tf.math.neg(input)))
+        return tf.math.add(positive, negative)
+    }
+
+    override fun toString(): String =
+        "PReLU(alphaInitializer=$alphaInitializer, alphaRegularizer=$alphaRegularizer, sharedAxes=$sharedAxes)"
+}

api/src/main/kotlin/org/jetbrains/kotlinx/dl/api/inference/keras/KerasConstants.kt

@@ -25,6 +25,7 @@ internal const val LAYER_BATCH_NORM: String = "BatchNormalization"
 internal const val LAYER_ACTIVATION: String = "Activation"
 internal const val LAYER_RELU: String = "ReLU"
 internal const val LAYER_ELU: String = "ELU"
+internal const val LAYER_PRELU: String = "PReLU"
 internal const val LAYER_LEAKY_RELU: String = "LeakyReLU"
 internal const val LAYER_THRESHOLDED_RELU = "ThresholdedReLU"
 internal const val LAYER_LSTM: String = "LSTM"

api/src/main/kotlin/org/jetbrains/kotlinx/dl/api/inference/keras/ModelLoader.kt

@@ -141,6 +141,7 @@ private fun convertToLayer(
         LAYER_ACTIVATION -> createActivationLayer(kerasLayer.config!!, kerasLayer.config.name!!)
         LAYER_RELU -> createReLULayer(kerasLayer.config!!, kerasLayer.config.name!!)
         LAYER_ELU -> createELULayer(kerasLayer.config!!, kerasLayer.config.name!!)
+        LAYER_PRELU -> createPReLULayer(kerasLayer.config!!, kerasLayer.config.name!!)
         LAYER_LEAKY_RELU -> createLeakyReLULayer(kerasLayer.config!!, kerasLayer.config.name!!)
         LAYER_THRESHOLDED_RELU -> createThresholdedReLULayer(kerasLayer.config!!, kerasLayer.config.name!!)
         LAYER_SOFTMAX -> createSoftmaxLayer(kerasLayer.config!!, kerasLayer.config.name!!)
@@ -390,6 +391,15 @@ private fun createELULayer(config: LayerConfig, name: String): Layer {
     )
 }
 
+private fun createPReLULayer(config: LayerConfig, name: String): Layer {
+    return PReLU(
+        alphaInitializer = convertToInitializer(config.alpha_initializer!!),
+        alphaRegularizer = convertToRegularizer(config.alpha_regularizer),
+        sharedAxes = config.shared_axes!!.toIntArray(),
+        name = name
+    )
+}
+
 private fun createLeakyReLULayer(config: LayerConfig, name: String): Layer {
     return LeakyReLU(
         alpha = config.alpha!!.toFloat(),

api/src/main/kotlin/org/jetbrains/kotlinx/dl/api/inference/keras/ModelSaver.kt

@@ -12,6 +12,7 @@ import org.jetbrains.kotlinx.dl.api.core.Sequential
 import org.jetbrains.kotlinx.dl.api.core.activation.Activations
 import org.jetbrains.kotlinx.dl.api.core.initializer.*
 import org.jetbrains.kotlinx.dl.api.core.layer.Layer
+import org.jetbrains.kotlinx.dl.api.core.layer.activation.PReLU
 import org.jetbrains.kotlinx.dl.api.core.layer.activation.LeakyReLU
 import org.jetbrains.kotlinx.dl.api.core.layer.activation.Softmax
 import org.jetbrains.kotlinx.dl.api.core.layer.activation.ThresholdedReLU
@@ -84,6 +85,7 @@ private fun convertToKerasLayer(layer: Layer, isKerasFullyCompatible: Boolean, i
         is Input -> createKerasInput(layer)
         is BatchNorm -> createKerasBatchNorm(layer, isKerasFullyCompatible)
         is ActivationLayer -> createKerasActivationLayer(layer)
+        is PReLU -> createKerasPReLULayer(layer, isKerasFullyCompatible)
         is LeakyReLU -> createKerasLeakyReLU(layer)
         is ThresholdedReLU -> createKerasThresholdedReLULayer(layer)
         is Add -> createKerasAddLayer(layer)
@@ -200,6 +202,17 @@ private fun createKerasActivationLayer(layer: ActivationLayer): KerasLayer {
     return KerasLayer(class_name = LAYER_ACTIVATION, config = configX)
 }
 
+private fun createKerasPReLULayer(layer: PReLU, isKerasFullyCompatible: Boolean): KerasLayer {
+    val configX = LayerConfig(
+        dtype = DATATYPE_FLOAT32,
+        alpha_initializer = convertToKerasInitializer(layer.alphaInitializer, isKerasFullyCompatible),
+        alpha_regularizer = convertToKerasRegularizer(layer.alphaRegularizer),
+        shared_axes = layer.sharedAxes?.toList(),
+        name = layer.name
+    )
+    return KerasLayer(class_name = LAYER_PRELU, config = configX)
+}
+
 private fun createKerasSoftmaxLayer(layer: Softmax): KerasLayer {
     val configX = LayerConfig(
         dtype = DATATYPE_FLOAT32,

api/src/main/kotlin/org/jetbrains/kotlinx/dl/api/inference/keras/config/LayerConfig.kt

@@ -15,6 +15,10 @@ internal data class LayerConfig(
     @Json(serializeNull = false)
     val alpha: Double? = null,
     @Json(serializeNull = false)
+    val alpha_initializer: KerasInitializer? = null,
+    @Json(serializeNull = false)
+    val alpha_regularizer: KerasRegularizer? = null,
+    @Json(serializeNull = false)
     val axis: Any? = null,
     @Json(serializeNull = false)
     var batch_input_shape: List<Int?>? = null,
@@ -119,6 +123,8 @@ internal data class LayerConfig(
     @Json(serializeNull = false)
     val scale: Any? = null,
     @Json(serializeNull = false)
+    val shared_axes: List<Int>? = null,
+    @Json(serializeNull = false)
     val sparse: Boolean? = null,
     @Json(serializeNull = false)
     val stateful: Boolean? = null,

api/src/test/kotlin/org/jetbrains/kotlinx/dl/api/core/layer/PReLUTest.kt

@@ -0,0 +1,188 @@
+/*
+ * Copyright 2020 JetBrains s.r.o. and Kotlin Deep Learning project contributors. All Rights Reserved.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlinx.dl.api.core.layer
+
+import org.jetbrains.kotlinx.dl.api.core.KGraph
+import org.jetbrains.kotlinx.dl.api.core.activation.EPS
+import org.jetbrains.kotlinx.dl.api.core.initializer.Constant
+import org.jetbrains.kotlinx.dl.api.core.initializer.RandomUniform
+import org.jetbrains.kotlinx.dl.api.core.layer.activation.PReLU
+import org.jetbrains.kotlinx.dl.api.core.shape.shapeFromDims
+import org.jetbrains.kotlinx.dl.api.core.shape.shapeOperand
+import org.junit.jupiter.api.Assertions
+import org.junit.jupiter.api.Test
+import org.tensorflow.Graph
+import org.tensorflow.Session
+import org.tensorflow.Shape
+import org.tensorflow.op.Ops
+
+class PReLUTest : ActivationLayerTest() {
+    @Test
+    fun default() {
+        val input = arrayOf(
+            floatArrayOf(-3.0f, -1.0f, 0.0f, 2.0f),
+            floatArrayOf(2.5f, 1.0f, -5.0f, 4.2f),
+        )
+        val expected = arrayOf(
+            floatArrayOf(0.0f, 0.0f, 0.0f, 2.0f),
+            floatArrayOf(2.5f, 1.0f, 0.0f, 4.2f),
+        )
+        val layer = PReLU()
+
+        Graph().use { graph ->
+            Session(graph).use { session ->
+                val tf = Ops.create(graph)
+                KGraph(graph.toGraphDef()).use { kGraph ->
+                    val inputShape = Shape.make(input.size.toLong(), input[0].size.toLong())
+                    layer.build(tf, kGraph, inputShape)
+
+                    val inputOp = tf.constant(input)
+                    val isTraining = tf.constant(false)
+                    val numberOfLosses = tf.constant(1.0f)
+                    val output = layer.forward(tf, inputOp, isTraining, numberOfLosses).asOutput()
+                    kGraph.initializeGraphVariables(session)
+                    val outputTensor = session.runner().fetch(output).run().first()
+
+                    // Check the output shape is correct.
+                    val expectedShape = inputShape
+                    val actualShape = shapeFromDims(*outputTensor.shape())
+                    Assertions.assertEquals(expectedShape, actualShape)
+
+                    // Check output is correct.
+                    val actual = Array(input.size) { FloatArray(input[0].size) }
+                    outputTensor.copyTo(actual)
+                    for (i in actual.indices) {
+                        Assertions.assertArrayEquals(
+                            expected[i],
+                            actual[i],
+                            EPS
+                        )
+                    }
+                }
+            }
+        }
+    }
+
+    @Test
+    fun withInitializer() {
+        val input = arrayOf(
+            floatArrayOf(-3.0f, -1.0f, 0.0f, 2.0f),
+            floatArrayOf(2.5f, 1.0f, -5.0f, 4.2f),
+        )
+        val expected = arrayOf(
+            floatArrayOf(-6.0f, -2.0f, 0.0f, 2.0f),
+            floatArrayOf(2.5f, 1.0f, -10.0f, 4.2f),
+        )
+        val layer = PReLU(Constant(2.0f))
+
+        Graph().use { graph ->
+            Session(graph).use { session ->
+                val tf = Ops.create(graph)
+                KGraph(graph.toGraphDef()).use { kGraph ->
+                    val inputShape = Shape.make(input.size.toLong(), input[0].size.toLong())
+                    layer.build(tf, kGraph, inputShape)
+
+                    val inputOp = tf.constant(input)
+                    val isTraining = tf.constant(false)
+                    val numberOfLosses = tf.constant(1.0f)
+                    val output = layer.forward(tf, inputOp, isTraining, numberOfLosses).asOutput()
+                    kGraph.initializeGraphVariables(session)
+                    val outputTensor = session.runner().fetch(output).run().first()
+
+                    // Check the output shape is correct.
+                    val expectedShape = inputShape
+                    val actualShape = shapeFromDims(*outputTensor.shape())
+                    Assertions.assertEquals(expectedShape, actualShape)
+
+                    // Check output is correct.
+                    val actual = Array(input.size) { FloatArray(input[0].size) }
+                    outputTensor.copyTo(actual)
+                    for (i in actual.indices) {
+                        Assertions.assertArrayEquals(
+                            expected[i],
+                            actual[i],
+                            EPS
+                        )
+                    }
+                }
+            }
+        }
+    }
+
+    @Test
+    fun withSharedAxes() {
+        val input = arrayOf(
+            arrayOf(
+                floatArrayOf(-3.0f, -1.0f, 0.0f, 2.0f),
+                floatArrayOf(2.5f, 1.0f, -5.0f, 4.2f),
+            ),
+            arrayOf(
+                floatArrayOf(2.0f, 0.5f, 1.0f, 4.0f),
+                floatArrayOf(-3.0f, -1.0f, 5.0f, -0.3f),
+            )
+        )
+        val initializerSeed = 42L
+        val layer = PReLU(RandomUniform(seed = initializerSeed), sharedAxes = intArrayOf(1))
+
+        Graph().use { graph ->
+            Session(graph).use { session ->
+                val tf = Ops.create(graph)
+                KGraph(graph.toGraphDef()).use { kGraph ->
+                    val inputShape = Shape.make(
+                        input.size.toLong(), input[0].size.toLong(), input[0][0].size.toLong()
+                    )
+                    layer.build(tf, kGraph, inputShape)
+
+                    val inputOp = tf.constant(input)
+                    val isTraining = tf.constant(false)
+                    val numberOfLosses = tf.constant(1.0f)
+                    val output = layer.forward(tf, inputOp, isTraining, numberOfLosses).asOutput()
+                    kGraph.initializeGraphVariables(session)
+                    val outputTensor = session.runner().fetch(output).run().first()
+
+                    // Check the output shape is correct.
+                    val expectedShape = inputShape
+                    val actualShape = shapeFromDims(*outputTensor.shape())
+                    Assertions.assertEquals(expectedShape, actualShape)
+
+                    // TODO: find a better way to verify the output is correct.
+                    // Check output is correct.
+                    val alphaInitializer = RandomUniform(seed = initializerSeed)
+                    val alphaTensor = session.runner().fetch(
+                        alphaInitializer.initialize(
+                            input[0][0].size,
+                            input[0][0].size,
+                            tf,
+                            shapeOperand(tf, Shape.make(1, input[0][0].size.toLong())),
+                            "temp_init"
+                        ).asOutput()
+                    ).run().first()
+                    val alpha = Array(1) { FloatArray(input[0][0].size) }
+                    alphaTensor.copyTo(alpha)
+
+                    val actual = Array(input.size) {
+                        Array(input[0].size) {
+                            FloatArray(input[0][0].size)
+                        }
+                    }
+                    outputTensor.copyTo(actual)
+                    for (i in actual.indices) {
+                        for (j in actual[i].indices) {
+                            for (k in actual[i][j].indices) {
+                                val it = input[i][j][k]
+                                Assertions.assertEquals(
+                                    if ( it < 0) it * alpha[0][k] else it,
+                                    actual[i][j][k],
+                                    EPS
+                                )
+                            }
+                        }
+                    }
+                }
+            }
+        }
+    }
+}