exp: Implement TensorFlow distribution strategies

12 files changed, 384 insertions, 27 deletions

diff --git a/opendc-experiments/opendc-experiments-tf20/src/main/kotlin/org/opendc/experiments/tf20/TensorFlowExperiment.kt b/opendc-experiments/opendc-experiments-tf20/src/main/kotlin/org/opendc/experiments/tf20/TensorFlowExperiment.kt
new file mode 100644
index 00000000..9a48aced
--- /dev/null
+++ b/opendc-experiments/opendc-experiments-tf20/src/main/kotlin/org/opendc/experiments/tf20/TensorFlowExperiment.kt
@@ -0,0 +1,72 @@
+/*
+ * Copyright (c) 2021 AtLarge Research
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+
+package org.opendc.experiments.tf20
+
+import io.opentelemetry.api.metrics.MeterProvider
+import io.opentelemetry.sdk.metrics.SdkMeterProvider
+import org.opendc.experiments.tf20.core.SimTFDevice
+import org.opendc.experiments.tf20.distribute.*
+import org.opendc.experiments.tf20.keras.AlexNet
+import org.opendc.experiments.tf20.util.MLEnvironmentReader
+import org.opendc.harness.dsl.Experiment
+import org.opendc.harness.dsl.anyOf
+import org.opendc.simulator.compute.power.LinearPowerModel
+import org.opendc.simulator.core.runBlockingSimulation
+import org.opendc.telemetry.sdk.toOtelClock
+
+/**
+ * Experiments with the TensorFlow simulation model.
+ */
+public class TensorFlowExperiment : Experiment(name = "tf20") {
+    /**
+     * The environment file to use.
+     */
+    private val environmentFile by anyOf("/kth.json")
+
+    /**
+     * The batch size used.
+     */
+    private val batchSize by anyOf(16, 32, 64, 128)
+
+    override fun doRun(repeat: Int): Unit = runBlockingSimulation {
+        val meterProvider: MeterProvider = SdkMeterProvider
+            .builder()
+            .setClock(clock.toOtelClock())
+            .build()
+        val meter = meterProvider.get("opendc-tf20")
+
+        val def = MLEnvironmentReader(TensorFlowExperiment::class.java.getResourceAsStream(environmentFile)).read().first()
+        val device = SimTFDevice(
+            def.uid, def.meta["gpu"] as Boolean, coroutineContext, clock, meter, def.model.cpus[0],
+            def.model.memory[0], LinearPowerModel(250.0, 60.0)
+        )
+        val strategy = OneDeviceStrategy(device)
+
+        val model = AlexNet(batchSize.toLong())
+        model.use {
+            it.compile(strategy)
+
+            it.fit(epochs = 9088 / batchSize, batchSize = batchSize)
+        }
+    }
+}
diff --git a/opendc-experiments/opendc-experiments-tf20/src/main/kotlin/org/opendc/experiments/tf20/distribute/MirroredStrategy.kt b/opendc-experiments/opendc-experiments-tf20/src/main/kotlin/org/opendc/experiments/tf20/distribute/MirroredStrategy.kt
new file mode 100644
index 00000000..8caa7ec9
--- /dev/null
+++ b/opendc-experiments/opendc-experiments-tf20/src/main/kotlin/org/opendc/experiments/tf20/distribute/MirroredStrategy.kt
@@ -0,0 +1,40 @@
+/*
+ * Copyright (c) 2021 AtLarge Research
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+
+package org.opendc.experiments.tf20.distribute
+
+import kotlinx.coroutines.coroutineScope
+import kotlinx.coroutines.launch
+import org.opendc.experiments.tf20.core.TFDevice
+
+/**
+ * A distribution [Strategy] that supports synchronous distributed training on multiple GPUs on one machine.
+ *
+ * It creates one replica per GPU device. Each variable in the model is mirrored across all the replicas.
+ */
+public class MirroredStrategy(val devices: List<TFDevice>) : Strategy {
+    override suspend fun run(forward: Double, backward: Double, batchSize: Int) = coroutineScope {
+        for (device in devices) {
+            launch { device.compute(forward * batchSize / devices.size + backward) }
+        }
+    }
+}
diff --git a/opendc-experiments/opendc-experiments-tf20/src/main/kotlin/org/opendc/experiments/tf20/distribute/OneDeviceStrategy.kt b/opendc-experiments/opendc-experiments-tf20/src/main/kotlin/org/opendc/experiments/tf20/distribute/OneDeviceStrategy.kt
new file mode 100644
index 00000000..271fab98
--- /dev/null
+++ b/opendc-experiments/opendc-experiments-tf20/src/main/kotlin/org/opendc/experiments/tf20/distribute/OneDeviceStrategy.kt
@@ -0,0 +1,34 @@
+/*
+ * Copyright (c) 2021 AtLarge Research
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+
+package org.opendc.experiments.tf20.distribute
+
+import org.opendc.experiments.tf20.core.TFDevice
+
+/**
+ * A distribution [Strategy] that places all variables and computation on a single specified device.
+ */
+public class OneDeviceStrategy(val device: TFDevice) : Strategy {
+    override suspend fun run(forward: Double, backward: Double, batchSize: Int) {
+        device.compute(forward * batchSize + backward)
+    }
+}
diff --git a/opendc-experiments/opendc-experiments-tf20/src/main/kotlin/org/opendc/experiments/tf20/distribute/Strategy.kt b/opendc-experiments/opendc-experiments-tf20/src/main/kotlin/org/opendc/experiments/tf20/distribute/Strategy.kt
new file mode 100644
index 00000000..5839c0df
--- /dev/null
+++ b/opendc-experiments/opendc-experiments-tf20/src/main/kotlin/org/opendc/experiments/tf20/distribute/Strategy.kt
@@ -0,0 +1,33 @@
+/*
+ * Copyright (c) 2021 AtLarge Research
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+
+package org.opendc.experiments.tf20.distribute
+
+/**
+ * A strategy for distributing TensorFlow state and computation over multiple devices.
+ */
+public interface Strategy {
+    /**
+     * Run the specified batch using the given strategy.
+     */
+    public suspend fun run(forward: Double, backward: Double, batchSize: Int)
+}
diff --git a/opendc-experiments/opendc-experiments-tf20/src/main/kotlin/org/opendc/experiments/tf20/keras/Sequential.kt b/opendc-experiments/opendc-experiments-tf20/src/main/kotlin/org/opendc/experiments/tf20/keras/Sequential.kt
index 3cc6b690..411ddb59 100644
--- a/opendc-experiments/opendc-experiments-tf20/src/main/kotlin/org/opendc/experiments/tf20/keras/Sequential.kt
+++ b/opendc-experiments/opendc-experiments-tf20/src/main/kotlin/org/opendc/experiments/tf20/keras/Sequential.kt
@@ -24,30 +24,35 @@ package org.opendc.experiments.tf20.keras
 
 import org.opendc.experiments.tf20.keras.layer.Layer
 import org.opendc.experiments.tf20.keras.layer.core.Input
+import org.opendc.experiments.tf20.keras.shape.TensorShape
 
 /**
- * Sequential model groups a linear stack of layers into a TensorFlow Model.
+ * Sequential model groups a linear stack of layers into a TensorFlow TrainableModel.
  *
  * @param [layers] The layers to describe the model design.
  */
-public class Sequential(vararg layers: Layer) : TrainableModel() {
+public class Sequential(vararg layers: Layer) : TrainableModel(*layers) {
+    override fun buildLayers() {
+        val inputShape = TensorShape(*inputLayer.packedDims)
+        inputLayer.inputTensor = inputShape
+        inputLayer.build(inputShape)
+        var nextShape: TensorShape = inputLayer.getOutputShape(inputShape)
+        inputLayer.outputTensor = nextShape
 
-    /**
-     * The layers to describe the model design. Main part of the internal state of the model.
-     */
-    public val layers: List<Layer> = listOf(*layers)
+        layers.filter { it !is Input }.forEach {
+            it.inputTensor = nextShape
+            it.build(nextShape)
 
-    /**
-     * First layer that is responsible for the input shape of the Neural Network.
-     */
-    public val inputLayer: Input
-        get() = layers[0] as Input
+            nextShape = it.getOutputShape(nextShape)
+            it.outputTensor = nextShape
+        }
+    }
 
-    /**
-     * Returns input dimensions in order HWC (height, width, channels)
-     */
-    public val inputDimensions: LongArray
-        get() = (layers[0] as Input).packedDims
+    override fun forward(): Double {
+        return layers.sumByDouble { it.forward() }
+    }
 
-    override fun close() {}
+    override fun backward(): Double {
+        return layers.sumByDouble { it.backward() }
+    }
 }
diff --git a/opendc-experiments/opendc-experiments-tf20/src/main/kotlin/org/opendc/experiments/tf20/keras/TrainableModel.kt b/opendc-experiments/opendc-experiments-tf20/src/main/kotlin/org/opendc/experiments/tf20/keras/TrainableModel.kt
index 421dec73..2cac6cbc 100644
--- a/opendc-experiments/opendc-experiments-tf20/src/main/kotlin/org/opendc/experiments/tf20/keras/TrainableModel.kt
+++ b/opendc-experiments/opendc-experiments-tf20/src/main/kotlin/org/opendc/experiments/tf20/keras/TrainableModel.kt
@@ -22,10 +22,107 @@
 
 package org.opendc.experiments.tf20.keras
 
+import org.opendc.experiments.tf20.distribute.Strategy
+import org.opendc.experiments.tf20.keras.layer.Layer
+import org.opendc.experiments.tf20.keras.layer.core.Input
+
 /**
- * Base abstract class for all trainable models.
+ * A model groups layers into an object with training and inference features.
  */
-public abstract class TrainableModel : AutoCloseable {
+public abstract class TrainableModel(vararg layers: Layer) : AutoCloseable {
+    /**
+     * The layers to describe the model design. Main part of the internal state of the model.
+     */
+    public val layers: List<Layer> = listOf(*layers)
+
+    /**
+     * First layer that is responsible for the input shape of the Neural Network.
+     */
+    public val inputLayer: Input
+        get() = layers[0] as Input
+
+    /**
+     * Returns input dimensions in order HWC (height, width, channels)
+     */
+    public val inputDimensions: LongArray
+        get() = (layers[0] as Input).packedDims
+
+    /**
+     * Layers indexed by name.
+     */
+    protected val layersByName: MutableMap<String, Layer> = mutableMapOf()
+
+    /**
+     * A flag to indicate that the model is compiled.
+     */
+    public var isCompiled: Boolean = false
+        private set
+
+    /**
+     * The strategy that is being used.
+     */
+    private lateinit var strategy: Strategy
+
+    /**
+     * Common method for building the initial part of the model static graph.
+     */
+    protected abstract fun buildLayers()
+
+    /**
+     * Perform a forward propagation.
+     */
+    protected abstract fun forward(): Double
+
+    /**
+     * Perform a backward propagation.
+     */
+    protected abstract fun backward(): Double
+
+    init {
+        for (layer in layers) {
+            if (layersByName.containsKey(layer.name)) {
+                throw IllegalArgumentException(layer.name)
+            } else {
+                layersByName[layer.name] = layer
+            }
+
+            layer.parentModel = this
+        }
+    }
+
+    /**
+     * Configures the model for training.
+     *
+     * @param strategy The distribution strategy for training.
+     */
+    public fun compile(strategy: Strategy) {
+        check(!isCompiled) { "Model is already compiled." }
+
+        buildLayers()
+
+        this.strategy = strategy
+        this.isCompiled = true
+    }
+
+    /**
+     * Train the model for a fixed number of [epochs] (iterations over a dataset).
+     *
+     * @param [epochs] Number of epochs to train the model. An epoch is an iteration over the entire x and y data provided.
+     * @param [batchSize] Number of samples per gradient update.
+     */
+    public suspend fun fit(epochs: Int = 5, batchSize: Int = 32) {
+        check(isCompiled) { "Model not yet compiled." }
+
+        val forwardFlops = forward()
+        val backwardFlops = backward()
+
+        for (i in 1..epochs) {
+            strategy.run(forwardFlops, backwardFlops, batchSize)
+        }
+    }
+
+    override fun close() {
+    }
 
     override fun toString(): String {
         return "TrainableModel ${super.toString()}"
diff --git a/opendc-experiments/opendc-experiments-tf20/src/main/kotlin/org/opendc/experiments/tf20/keras/layer/Layer.kt b/opendc-experiments/opendc-experiments-tf20/src/main/kotlin/org/opendc/experiments/tf20/keras/layer/Layer.kt
index eafdf63b..143b27f0 100644
--- a/opendc-experiments/opendc-experiments-tf20/src/main/kotlin/org/opendc/experiments/tf20/keras/layer/Layer.kt
+++ b/opendc-experiments/opendc-experiments-tf20/src/main/kotlin/org/opendc/experiments/tf20/keras/layer/Layer.kt
@@ -22,6 +22,7 @@
 
 package org.opendc.experiments.tf20.keras.layer
 
+import org.opendc.experiments.tf20.keras.TrainableModel
 import org.opendc.experiments.tf20.keras.shape.TensorShape
 
 /**
@@ -31,6 +32,23 @@ import org.opendc.experiments.tf20.keras.shape.TensorShape
  */
 public abstract class Layer(public val name: String) {
     /**
+     * TrainableModel in which the layer exists.
+     */
+    internal var parentModel: TrainableModel? = null
+
+    /**
+     * The input shape of the layer.
+     */
+    public lateinit var inputTensor: TensorShape
+        internal set
+
+    /**
+     * The output shape of the layer.
+     */
+    public lateinit var outputTensor: TensorShape
+        internal set
+
+    /**
      * Build the layer for the specified [inputShape].
      *
      * @param [inputShape] Input shape, result of [getOutputShape] call from previous layer.
@@ -41,4 +59,14 @@ public abstract class Layer(public val name: String) {
      * Compute output shape of this layer, based on [inputShape] and [Layer] type.
      */
     public abstract fun getOutputShape(inputShape: TensorShape): TensorShape
+
+    /**
+     * Perform a forward propagation
+     */
+    public abstract fun forward(): Double
+
+    /**
+     * Perform a backward propagation.
+     */
+    public abstract fun backward(): Double
 }
diff --git a/opendc-experiments/opendc-experiments-tf20/src/main/kotlin/org/opendc/experiments/tf20/keras/layer/conv/Conv2D.kt b/opendc-experiments/opendc-experiments-tf20/src/main/kotlin/org/opendc/experiments/tf20/keras/layer/conv/Conv2D.kt
index c8078bf8..bf4725b7 100644
--- a/opendc-experiments/opendc-experiments-tf20/src/main/kotlin/org/opendc/experiments/tf20/keras/layer/conv/Conv2D.kt
+++ b/opendc-experiments/opendc-experiments-tf20/src/main/kotlin/org/opendc/experiments/tf20/keras/layer/conv/Conv2D.kt
@@ -45,12 +45,10 @@ public class Conv2D(
     private var padHeight: Double = 0.0
     private var padWidth: Double = 0.0
 
-    override fun build(inputShape: TensorShape) {
-        TODO("not implemented")
-    }
+    override fun build(inputShape: TensorShape) {}
 
     override fun getOutputShape(inputShape: TensorShape): TensorShape {
-        check(filter[2] != inputShape[3]) { "Input channel ${filter[2]} and ${inputShape[3]} shall match" }
+        check(filter[2] == inputShape[3]) { "Input channel ${filter[2]} and ${inputShape[3]} shall match" }
 
         var outHeight = 0L
         var outWidth = 0L
@@ -74,6 +72,27 @@ public class Conv2D(
         return TensorShape(inputShape[0], outHeight, outWidth, filter[3])
     }
 
+    override fun forward(): Double {
+        // Mul and add per output pixel: kernel_w x kernel_h x in_channel
+        var flops: Long = (2 * filter[0] * filter[1] * filter[2])
+
+        val output = outputTensor
+        // Flops per output map.
+        flops *= output[1] * output[2] * filter[3]
+
+        // Flops across multiple input patches.
+        flops *= inputTensor[0]
+
+        if (activation == Activation.Relu) {
+            flops += output[0] * output[1] * output[2] * output[3]
+        }
+
+        // return paramsNum() * output.H * output.W * FLOAT_BYTES / MILLION
+        return flops * 4.0 / 1_000_000
+    }
+
+    override fun backward(): Double = forward()
+
     override fun toString(): String {
         return "Conv2D[filter=${filter.contentToString()}, strides=${strides.contentToString()}, activation=$activation, padding=$padding]"
     }
diff --git a/opendc-experiments/opendc-experiments-tf20/src/main/kotlin/org/opendc/experiments/tf20/keras/layer/core/ActivationLayer.kt b/opendc-experiments/opendc-experiments-tf20/src/main/kotlin/org/opendc/experiments/tf20/keras/layer/core/ActivationLayer.kt
index befcfe50..fd25cea6 100644
--- a/opendc-experiments/opendc-experiments-tf20/src/main/kotlin/org/opendc/experiments/tf20/keras/layer/core/ActivationLayer.kt
+++ b/opendc-experiments/opendc-experiments-tf20/src/main/kotlin/org/opendc/experiments/tf20/keras/layer/core/ActivationLayer.kt
@@ -40,6 +40,10 @@ public class ActivationLayer(
 
     override fun getOutputShape(inputShape: TensorShape): TensorShape = inputShape
 
+    override fun forward(): Double = 0.0
+
+    override fun backward(): Double = forward()
+
     override fun toString(): String {
         return "ActivationLayer[activation=$activation]"
     }
diff --git a/opendc-experiments/opendc-experiments-tf20/src/main/kotlin/org/opendc/experiments/tf20/keras/layer/core/Input.kt b/opendc-experiments/opendc-experiments-tf20/src/main/kotlin/org/opendc/experiments/tf20/keras/layer/core/Input.kt
index b56284d6..6619ccc0 100644
--- a/opendc-experiments/opendc-experiments-tf20/src/main/kotlin/org/opendc/experiments/tf20/keras/layer/core/Input.kt
+++ b/opendc-experiments/opendc-experiments-tf20/src/main/kotlin/org/opendc/experiments/tf20/keras/layer/core/Input.kt
@@ -32,16 +32,18 @@ public class Input(vararg dims: Long, name: String) : Layer(name) {
     /**
      * Input data dimensions. Rank = 3 or 4 for most popular supported cases.
      */
-    public var packedDims: LongArray = dims
+    public val packedDims: LongArray = dims
 
-    override fun build(inputShape: TensorShape) {
-        TODO("not implemented")
-    }
+    override fun build(inputShape: TensorShape) {}
 
     override fun getOutputShape(inputShape: TensorShape): TensorShape {
         return inputShape
     }
 
+    override fun forward(): Double = 0.0
+
+    override fun backward(): Double = 0.0
+
     override fun toString(): String {
         return "Input[shape=${packedDims.contentToString()}]"
     }
diff --git a/opendc-experiments/opendc-experiments-tf20/src/main/kotlin/org/opendc/experiments/tf20/keras/layer/pool/Pool2D.kt b/opendc-experiments/opendc-experiments-tf20/src/main/kotlin/org/opendc/experiments/tf20/keras/layer/pool/Pool2D.kt
index b6ec78e0..3c6b15bb 100644
--- a/opendc-experiments/opendc-experiments-tf20/src/main/kotlin/org/opendc/experiments/tf20/keras/layer/pool/Pool2D.kt
+++ b/opendc-experiments/opendc-experiments-tf20/src/main/kotlin/org/opendc/experiments/tf20/keras/layer/pool/Pool2D.kt
@@ -71,6 +71,22 @@ public class Pool2D(
         return TensorShape(inputShape[0], outHeight, outWidth, inputShape[3])
     }
 
+    override fun forward(): Double {
+        val output = outputTensor
+        // Per output pixel: kernel_w x kernel_h x in_channel
+        var flops: Long = 2 * poolSize[1] * poolSize[2] * inputTensor[3]
+
+        // Flops per output map.
+        flops *= output[2] * output[1]
+
+        // Flops across multiple input patches.
+        flops *= inputTensor[0]
+
+        return flops * 4.0 / 1_000_000
+    }
+
+    override fun backward(): Double = forward()
+
     override fun toString(): String {
         return "MaxPool2D[poolSize=${poolSize.contentToString()}, strides=${strides.contentToString()}, padding=$padding]"
     }
diff --git a/opendc-experiments/opendc-experiments-tf20/src/main/kotlin/org/opendc/experiments/tf20/keras/layer/regularization/Dropout.kt b/opendc-experiments/opendc-experiments-tf20/src/main/kotlin/org/opendc/experiments/tf20/keras/layer/regularization/Dropout.kt
index 1fead435..ff5f7711 100644
--- a/opendc-experiments/opendc-experiments-tf20/src/main/kotlin/org/opendc/experiments/tf20/keras/layer/regularization/Dropout.kt
+++ b/opendc-experiments/opendc-experiments-tf20/src/main/kotlin/org/opendc/experiments/tf20/keras/layer/regularization/Dropout.kt
@@ -46,5 +46,12 @@ public class Dropout(
         return inputShape
     }
 
+    override fun forward(): Double {
+        val output = outputTensor
+        return output[0] * output[1] * output[2] * output[3] * 4.0 / 1_000_000
+    }
+
+    override fun backward(): Double = forward()
+
     override fun toString(): String = "Dropout[keepProbability=$keepProbability]"
 }