diff --git a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/nn/DropoutMaskingTest.kt b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/nn/DropoutMaskingTest.kt new file mode 100644 index 00000000..93a604ce --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/nn/DropoutMaskingTest.kt @@ -0,0 +1,62 @@ +package sk.ainet.exec.nn + +import kotlin.math.abs +import kotlin.random.Random +import kotlin.test.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertTrue +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.context.Phase +import sk.ainet.lang.nn.layers.Dropout +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.types.FP32 + +/** + * Value-level dropout behavior on the real CPU backend (the lang-core default + * context uses the shape-only Void backend, so these assertions live here). + */ +class DropoutMaskingTest { + + private val trainCtx = DirectCpuExecutionContext(phase = Phase.TRAIN) + private val evalCtx = DirectCpuExecutionContext(phase = Phase.EVAL) + + private fun ones(n: Int, ctx: DirectCpuExecutionContext): Tensor = + ctx.fromFloatArray(Shape(n), FP32::class, FloatArray(n) { 1f }) + + @Test + fun training_phase_zeroes_or_scales_every_element() { + val layer = Dropout(p = 0.5f, random = Random(42)) + val out = layer.forward(ones(64, trainCtx), trainCtx).data.copyToFloatArray() + + assertTrue(out.all { it == 0f || it == 2f }, "inverted dropout must zero or scale by 1/(1-p)") + assertTrue(out.any { it == 0f }, "expected at least one dropped element") + assertTrue(out.any { it == 2f }, "expected at least one surviving element") + } + + @Test + fun expected_activation_is_preserved() { + val layer = Dropout(p = 0.3f, random = Random(123)) + val out = layer.forward(ones(10_000, trainCtx), trainCtx).data.copyToFloatArray() + + val mean = out.average().toFloat() + assertTrue(abs(mean - 1f) < 0.05f, "inverted dropout should keep the mean ~1, was $mean") + } + + @Test + fun seeded_random_is_reproducible() { + val x = trainCtx.fromFloatArray(Shape(16), FP32::class, FloatArray(16) { it.toFloat() }) + + val y1 = Dropout(p = 0.5f, random = Random(7)).forward(x, trainCtx) + val y2 = Dropout(p = 0.5f, random = Random(7)).forward(x, trainCtx) + + assertContentEquals(y1.data.copyToFloatArray(), y2.data.copyToFloatArray()) + } + + @Test + fun inference_phase_is_exact_identity() { + val x = evalCtx.fromFloatArray(Shape(4), FP32::class, floatArrayOf(1f, 2f, 3f, 4f)) + val y = Dropout(p = 0.5f, random = Random(1)).forward(x, evalCtx) + assertContentEquals(floatArrayOf(1f, 2f, 3f, 4f), y.data.copyToFloatArray()) + } +} diff --git a/skainet-lang/skainet-lang-core/api/android/skainet-lang-core.api b/skainet-lang/skainet-lang-core/api/android/skainet-lang-core.api index c9c152e6..a68486ca 100644 --- a/skainet-lang/skainet-lang-core/api/android/skainet-lang-core.api +++ b/skainet-lang/skainet-lang-core/api/android/skainet-lang-core.api @@ -1641,8 +1641,8 @@ public final class sk/ainet/lang/nn/hooks/TapeRecorder$Entry { public final class sk/ainet/lang/nn/layers/Dropout : sk/ainet/lang/nn/Module { public fun ()V - public fun (FZLjava/lang/String;)V - public synthetic fun (FZLjava/lang/String;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public fun (FZLjava/lang/String;Lkotlin/random/Random;)V + public synthetic fun (FZLjava/lang/String;Lkotlin/random/Random;ILkotlin/jvm/internal/DefaultConstructorMarker;)V public fun forward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/context/ExecutionContext;)Lsk/ainet/lang/tensor/Tensor; public fun getModules ()Ljava/util/List; public fun getName ()Ljava/lang/String; diff --git a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api b/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api index 4c7cbad8..b631e303 100644 --- a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api +++ b/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api @@ -1930,8 +1930,8 @@ public final class sk/ainet/lang/nn/hooks/TapeRecorder$Entry { public final class sk/ainet/lang/nn/layers/Dropout : sk/ainet/lang/nn/Module { public fun ()V - public fun (FZLjava/lang/String;)V - public synthetic fun (FZLjava/lang/String;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public fun (FZLjava/lang/String;Lkotlin/random/Random;)V + public synthetic fun (FZLjava/lang/String;Lkotlin/random/Random;ILkotlin/jvm/internal/DefaultConstructorMarker;)V public fun forward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/context/ExecutionContext;)Lsk/ainet/lang/tensor/Tensor; public fun getModules ()Ljava/util/List; public fun getName ()Ljava/lang/String; diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/layers/Dropout.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/layers/Dropout.kt index 3b23ee42..3fb59683 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/layers/Dropout.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/layers/Dropout.kt @@ -3,19 +3,37 @@ package sk.ainet.lang.nn.layers import sk.ainet.context.ExecutionContext import sk.ainet.lang.nn.Module import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.data.TensorData +import sk.ainet.lang.types.BF16 import sk.ainet.lang.types.DType +import sk.ainet.lang.types.FP16 +import sk.ainet.lang.types.FP32 +import sk.ainet.lang.types.FP64 +import kotlin.random.Random /** - * Dropout layer that is aware of ExecutionContext phases. + * Dropout layer with inverted-dropout semantics, aware of ExecutionContext phases. * - * Current implementation keeps identity semantics on both TRAIN and EVAL phases, - * but exposes a context-aware forward that can later be extended to apply - * stochastic masking when ctx.inTraining is true. + * During training ([ExecutionContext.inTraining] and [training] both true, `p > 0`), + * each element is zeroed with probability [p] and the surviving elements are + * scaled by `1 / (1 - p)` so the expected activation stays constant. In the + * evaluation phase (or with [training] disabled) the layer is the identity — + * no rescaling is needed at inference time. + * + * The mask is a constant tensor combined with an element-wise multiply, so + * gradients flow to the surviving inputs only — the standard dropout backward — + * without needing a dedicated autograd rule. + * + * @param p probability of zeroing an element, in `[0, 1)` + * @param training manual override; set to false to disable masking regardless of phase + * @param name name of the module + * @param random RNG for the mask — inject a seeded [Random] for reproducible runs */ public class Dropout( public val p: Float = 0.5f, public var training: Boolean = true, - override val name: String = "Dropout" + override val name: String = "Dropout", + private val random: Random = Random.Default, ) : Module() { init { @@ -27,12 +45,31 @@ public class Dropout( get() = emptyList() /** - * Context-aware forward that can use ExecutionContext.phase. Hooks are dispatched if available. + * Context-aware forward: stochastic masking when ctx is in the training + * phase, identity otherwise. Hooks are dispatched if available. */ override fun forward(input: Tensor, ctx: ExecutionContext): Tensor = sk.ainet.lang.nn.hooks.withForwardHooks(ctx, this, input) { - // Placeholder behavior: identity in both phases. When RNG and elementwise ops are available, - // implement: if (ctx.inTraining && p > 0f) then output = input * mask / (1 - p) - input + if (!ctx.inTraining || !training || p == 0f) { + input + } else { + val mask = ctx.fromData(bernoulliMask(input, ctx), input.dtype) + ctx.ops.multiply(input, mask) + } } + + private fun bernoulliMask(input: Tensor, ctx: ExecutionContext): TensorData { + val scale = 1f / (1f - p) + return ctx.tensorDataFactory.init(input.shape, input.dtype) { + val keep = random.nextFloat() >= p + @Suppress("UNCHECKED_CAST") + when (input.dtype) { + FP32::class, FP16::class, BF16::class -> (if (keep) scale else 0f) as V + FP64::class -> (if (keep) scale.toDouble() else 0.0) as V + else -> throw UnsupportedOperationException( + "Dropout($name): unsupported dtype ${input.dtype} — floating-point tensors only" + ) + } + } + } } diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/nn/layers/DropoutPhaseTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/nn/layers/DropoutPhaseTest.kt index 3841ea7b..a081b586 100644 --- a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/nn/layers/DropoutPhaseTest.kt +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/nn/layers/DropoutPhaseTest.kt @@ -17,13 +17,15 @@ class DropoutPhaseTest { val layer = Dropout(p = 0.3f) val y = eval(base) { ctx -> layer.forward(x, ctx) } assertEquals(x.shape, y.shape) + kotlin.test.assertContentEquals(floatArrayOf(1f, 2f, 3f, 4f), y.data.copyToFloatArray()) } @Test - fun dropout_identity_in_train_phase_with_ctx() { + fun dropout_identity_in_train_phase_when_p_is_zero() { val x = base.fromFloatArray(Shape(2, 2), FP32::class, floatArrayOf(1f, 2f, 3f, 4f)) val layer = Dropout(p = 0.0f) val y = train(base) { ctx -> layer.forward(x, ctx) } assertEquals(x.shape, y.shape) + kotlin.test.assertContentEquals(floatArrayOf(1f, 2f, 3f, 4f), y.data.copyToFloatArray()) } } diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/nn/layers/DropoutTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/nn/layers/DropoutTest.kt index bedd8120..801c3c38 100644 --- a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/nn/layers/DropoutTest.kt +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/nn/layers/DropoutTest.kt @@ -1,12 +1,20 @@ package sk.ainet.lang.nn.layers +import kotlin.random.Random import kotlin.test.Test +import kotlin.test.assertContentEquals import kotlin.test.assertEquals +import sk.ainet.context.train +import sk.ainet.lang.nn.DefaultNeuralNetworkExecutionContext +import sk.ainet.lang.nn.NeuralNetworkExecutionContext import sk.ainet.lang.tensor.Shape import sk.ainet.lang.types.FP32 -import sk.ainet.lang.nn.NeuralNetworkExecutionContext -import sk.ainet.lang.nn.DefaultNeuralNetworkExecutionContext +/** + * Identity paths and shape behavior. The default context here uses the + * shape-only Void backend, so value-level masking behavior is covered in + * skainet-backend-cpu's DropoutMaskingTest. + */ class DropoutTest { private val ctx: NeuralNetworkExecutionContext = DefaultNeuralNetworkExecutionContext() @@ -15,8 +23,8 @@ class DropoutTest { fun p_zero_is_identity_in_training() { val x = ctx.fromFloatArray(Shape(3), FP32::class, floatArrayOf(1f, 2f, 3f)) val layer = Dropout(p = 0f, training = true) - val y = layer.forward(x, ctx) - assertEquals(x.shape, y.shape) + val y = train(ctx) { trainCtx -> layer.forward(x, trainCtx) } + assertContentEquals(floatArrayOf(1f, 2f, 3f), y.data.copyToFloatArray()) } @Test @@ -24,6 +32,14 @@ class DropoutTest { val x = ctx.fromFloatArray(Shape(2, 2), FP32::class, floatArrayOf(1f, 2f, 3f, 4f)) val layer = Dropout(p = 0.5f, training = false) val y = layer.forward(x, ctx) + assertContentEquals(floatArrayOf(1f, 2f, 3f, 4f), y.data.copyToFloatArray()) + } + + @Test + fun masked_path_preserves_the_shape() { + val x = ctx.fromFloatArray(Shape(2, 3, 4), FP32::class, FloatArray(24) { 1f }) + val layer = Dropout(p = 0.5f, random = Random(1)) + val y = train(ctx) { trainCtx -> layer.forward(x, trainCtx) } assertEquals(x.shape, y.shape) } }