From 043925303c070a29b993eec942f001551f1d5ba4 Mon Sep 17 00:00:00 2001 From: Marcel Luethi Date: Sat, 1 Aug 2026 16:51:02 +0200 Subject: [PATCH] Rework optimizer + LRSchedule * move tensor tree related methods to separate package * Rework optimizers to get the VType passed explicitly with .of methods. Add a LearningRateSchedule optimizer. * Update docs to new GradientOptimizer API --- AGENTS.md | 10 +- .../dimwit/optimizer/GradientOptimizer.scala | 156 +++++++++++------- .../main/scala/dimwit/tensor/TensorOps.scala | 3 +- .../scala/dimwit/tensortree/FloatTree.scala | 21 +++ .../optimizer/GradientOptimizerSuite.scala | 56 +++++-- docs/quickstart.md | 7 +- .../dimwit/basic/LogisticRegression.scala | 2 +- .../complex/VariationalAutoencoder.scala | 23 +-- mdocs/AGENTS.md | 10 +- mdocs/docs/quickstart.md | 3 +- 10 files changed, 190 insertions(+), 101 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index eae57c1..89c94f0 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -870,6 +870,7 @@ println(s"Block Hessian shapes: ${h_x1x1.shape}, ${h_x1x2.shape}, ${h_x2x1.shape ```scala import dimwit.* +import dimwit.Conversions.given import dimwit.optimizer.{GradientDescent, GradientOptimizer} import dimwit.random.Random @@ -903,7 +904,7 @@ val lossFunc = mse(trainData, trainLabels) val gradFunc = Autodiff.grad(lossFunc) // Create optimizer -val optimizer = GradientDescent(learningRate = 0.01f) +val optimizer = GradientDescent.of(VType[Float32])(learningRate = 0.01f) // Training loop with iterator val trained = optimizer.iterate(initModelParams)(gradFunc) @@ -918,9 +919,10 @@ val trained = optimizer.iterate(initModelParams)(gradFunc) ```scala import dimwit.optimizer.Lion +import dimwit.Conversions.given // enables implicit conversion from Float to Tensor[V] // Lion optimizer with momentum -val lionOptimizer = Lion(learningRate = 1e-3f, beta1 = 0.9f, beta2 = 0.99f, weightDecay = 0.0f) +val lionOptimizer = Lion.of(VType[Float32])(learningRate = 1e-3f, beta1 = 0.9f, beta2 = 0.99f, weightDecay = 0.0f) // Training with Lion val trainedLion = lionOptimizer.iterate(initModelParams)(gradFunc) @@ -932,6 +934,8 @@ val trainedLion = lionOptimizer.iterate(initModelParams)(gradFunc) ### Complete Training Example: Linear Regression ```scala +import dimwit.Conversions.given // enables implicit conversion from Float to Tensor[V] + // Define problem dimensions trait Sample derives Label trait InputDim derives Label @@ -963,7 +967,7 @@ val initRegressionParams = RegressionParams(initSlope, initIntercept) // Train val regressionGrad = Autodiff.grad(regressionLoss(xData, yData)) -val gdOptimizer = GradientDescent(learningRate = 0.1f) +val gdOptimizer = GradientDescent.of(VType[Float32])(learningRate = 0.1f) val finalParams = gdOptimizer.iterate(initRegressionParams)(regressionGrad) .take(100) diff --git a/core/src/main/scala/dimwit/optimizer/GradientOptimizer.scala b/core/src/main/scala/dimwit/optimizer/GradientOptimizer.scala index 9915b59..26aff0c 100644 --- a/core/src/main/scala/dimwit/optimizer/GradientOptimizer.scala +++ b/core/src/main/scala/dimwit/optimizer/GradientOptimizer.scala @@ -1,12 +1,11 @@ package dimwit.optimizer import dimwit.* - import dimwit.Conversions.given +import dimwit.autodiff.* +import dimwit.autodiff.Grad import dimwit.tensortree.* -import dimwit.tensortree.FloatTree.* import dimwit.tensortree.FloatTree.ops.* -import dimwit.autodiff.* /** Gradient optimizer interface with functional state management. * @@ -27,85 +26,96 @@ import dimwit.autodiff.* * optimizer.update(grads, params, state) * }}} */ -trait GradientOptimizer: - type State[_, V] +trait GradientOptimizer[V: IsFloating, State0[_]]: + + type State[P] = State0[P] // Core API - def init[V, Params: TensorTree: FloatTreeFor[V]](params: Params)(using IsFloating[V]): State[Params, V] - def update[V, Params: TensorTree: FloatTreeFor[V]](gradients: Grad[Params], params: Params, state: State[Params, V])(using IsFloating[V]): (Params, State[Params, V]) + def init[P](params: P)(using TensorTree[P], FloatTree[P, V]): State[P] + def update[P](gradients: Grad[P], params: P, state: State[P])(using TensorTree[P], FloatTree[P, V]): (P, State[P]) // Convenience: iterator with fixed gradient function - def iterateWithState[V, Params: TensorTree: FloatTreeFor[V]](init: Params)(df: Params => Grad[Params])(using IsFloating[V]): Iterator[(Params, State[Params, V])] = + def iterateWithState[P](init: P)(df: P => Grad[P])(using TensorTree[P], FloatTree[P, V]): Iterator[(P, State[P])] = Iterator.iterate((init, this.init(init))): (params, state) => val grads = df(params) update(grads, params, state) - def iterate[V, Params: TensorTree: FloatTreeFor[V]](init: Params)(df: Params => Grad[Params])(using IsFloating[V]): Iterator[Params] = + def iterate[P](init: P)(df: P => Grad[P])(using TensorTree[P], FloatTree[P, V]): Iterator[P] = iterateWithState(init)(df).map(_._1) -case class GradientDescent(learningRate: Double) extends GradientOptimizer: +object GradientDescent: + + def of[V: IsFloating](vtype: VType[V])(learningRate: Tensor0[V]): GradientDescent[V] = new GradientDescent(learningRate) + +type GradientDescentState[P] = Unit // empty state - type State[P, V] = Unit // Stateless optimizer +class GradientDescent[V: IsFloating](val learningRate: Tensor0[V]) extends GradientOptimizer[V, GradientDescentState]: - def init[V, Params: TensorTree: FloatTreeFor[V]](params: Params)(using IsFloating[V]): Unit = () + def init[P](params: P)(using TensorTree[P], FloatTree[P, V]): Unit = () - def update[V, Params: TensorTree: FloatTreeFor[V]](gradients: Grad[Params], params: Params, state: Unit)(using IsFloating[V]): (Params, Unit) = + def update[P](gradients: Grad[P], params: P, state: Unit)(using TensorTree[P], FloatTree[P, V]): (P, Unit) = val newParams = params -- gradients.value.scale(learningRate) (newParams, ()) -case class Lion(learningRate: Double, weightDecay: Double = 0.0f, beta1: Double = 0.9f, beta2: Double = 0.99f) extends GradientOptimizer: +case class LionState[P]( + momentums: P, + step: Tensor0[Int32] +) + +object Lion: + + def of[V](vtype: VType[V])(using IsFloating[V])(learningRate: Tensor0[V], weightDecay: Tensor0[V] = Tensor0(vtype)(0.0), beta1: Tensor0[V] = Tensor0(vtype)(0.9), beta2: Tensor0[V] = Tensor0(vtype)(0.99)): Lion[V] = new Lion(learningRate, weightDecay, beta1, beta2) - type State[P, V] = P // momentum state has same structure as params +class Lion[V: IsFloating](val learningRate: Tensor0[V], val weightDecay: Tensor0[V] = Tensor0(0.0), val beta1: Tensor0[V] = Tensor0(0.9), val beta2: Tensor0[V] = Tensor0(0.99)) extends GradientOptimizer[V, LionState]: - def init[V, Params: TensorTree: FloatTreeFor[V]](params: Params)(using IsFloating[V]): Params = - params.map([T <: Tuple] => - (n: Labels[T]) ?=> - (t: Tensor[T, V]) => - Tensor(t.shape).fill(0f) - ) + def init[P](params: P)(using TensorTree[P], FloatTree[P, V]): LionState[P] = + LionState(params.fillCopy(0f), step = 1) - def update[V, Params: TensorTree: FloatTreeFor[V]](gradients: Grad[Params], params: Params, momentums: Params)(using IsFloating[V]): (Params, Params) = + def update[P](gradients: Grad[P], params: P, state: LionState[P])(using TensorTree[P], FloatTree[P, V]): (P, LionState[P]) = // the direction (1 or -1) // is determined by the sign of the momentum + gradient - val updateDirection = (momentums **! beta1 ++ gradients.value **! (1f - beta1)).sign + val updateDirection = (state.momentums **! beta1 ++ gradients.value **! (1f - beta1)).sign val updatedParams = params -- updateDirection.scale(learningRate) -- params.scale(weightDecay) - val newMomentums = momentums **! beta2 ++ gradients.value **! (1f - beta2) + val newMomentums = state.momentums **! beta2 ++ gradients.value **! (1f - beta2) - (updatedParams, newMomentums) + (updatedParams, LionState(newMomentums, state.step + 1)) -case class AdamState[P, V: IsFloating]( +case class AdamState[P, V]( momentums: P, // momentums velocities: P, // velocities b1: Tensor0[V], // decay rate for momentums mᵗ b2: Tensor0[V] // decay rate for velocities vᵗ ) +object Adam: + + def of[V](vtype: VType[V])(using IsFloating[V])(learningRate: Tensor0[V], b1: Tensor0[V] = Tensor0(vtype)(0.9), b2: Tensor0[V] = Tensor0(vtype)(0.999), epsilon: Tensor0[V] = Tensor0(vtype)(1e-8)): Adam[V] = new Adam(learningRate, b1, b2, epsilon) + /** Implements the Adam optimization algorithm. * * @see [[https://arxiv.org/abs/1412.6980 Adam: A Method for Stochastic Optimization]] */ -case class Adam( - learningRate: Double, // step size (learning rate) - b1: Double = 0.9, // decay rate for momentums mᵗ - b2: Double = 0.999, // decay rate for velocities vᵗ - epsilon: Double = 1e-8 // small constant to prevent division by zero -) extends GradientOptimizer: +class Adam[V: IsFloating]( + val learningRate: Tensor0[V], + b1: Tensor0[V] = Tensor0(0.9), // decay rate for momentums mᵗ + b2: Tensor0[V] = Tensor0(0.999), // decay rate for velocities vᵗ + epsilon: Tensor0[V] = Tensor0(1e-8) // small constant to prevent division by zero +) extends GradientOptimizer[V, [P] =>> AdamState[P, V]]: private val β1 = b1 private val β2 = b2 + private val ε = epsilon - type State[P, V] = AdamState[P, V] - - def init[V, Params: TensorTree: FloatTreeFor[V]](params: Params)(using IsFloating[V]): State[Params, V] = + def init[P](params: P)(using TensorTree[P], FloatTree[P, V]): AdamState[P, V] = def zeros = params.fillCopy(0f) - AdamState[Params, V](zeros, zeros, b1 = Tensor0(VType[V])(1f), b2 = Tensor0(VType[V])(1f)) + AdamState(zeros, zeros, b1 = Tensor0(VType[V])(1f), b2 = Tensor0(VType[V])(1f)) - def update[V, Params: TensorTree: FloatTreeFor[V]]( - gradients: Grad[Params], - params: Params, - state: State[Params, V] - )(using IsFloating[V]): (Params, State[Params, V]) = + def update[P]( + gradients: Grad[P], + params: P, + state: AdamState[P, V] + )(using TensorTree[P], FloatTree[P, V]): (P, AdamState[P, V]) = // rename state variables to last time step for clarity val `mₜ₋₁` = state.momentums val `vₜ₋₁` = state.velocities @@ -114,17 +124,15 @@ case class Adam( // rename parameters for internal clarity val α = learningRate - val ε = epsilon - val `θₜ₋₁` = params - // update moments for bias correction - val β1ₜ = `β1ₜ₋₁` * β1 - val β2ₜ = `β2ₜ₋₁` * β2 + val `θₜ₋₁` = params // Adam implementation val gₜ = gradients.value val mᵗ = `β1` **! `mₜ₋₁` ++ (1f - `β1`) **! gₜ val vᵗ = `β2` **! `vₜ₋₁` ++ (1f - `β2`) **! gₜ.pow(2) + val β1ₜ = `β1ₜ₋₁` * β1 + val β2ₜ = `β2ₜ₋₁` * β2 val m̂ = mᵗ `//!` (1f - `β1ₜ`) val v̂ = vᵗ `//!` (1f - `β2ₜ`) val θₜ = `θₜ₋₁` -- (α **! m̂) `//` (v̂.sqrt ++! ε) @@ -141,20 +149,18 @@ case class Adam( * @param learningRate The step size. * @param weightDecayFactor The coefficient for weight decay (lambda). */ -case class AdamW( - val adam: Adam, - val weightDecayFactor: Double -) extends GradientOptimizer: - - type State[P, V] = adam.State[P, V] - - def init[V, Params: TensorTree: FloatTreeFor[V]](params: Params)(using IsFloating[V]): State[Params, V] = adam.init(params) - - def update[V, Params: TensorTree: FloatTreeFor[V]]( - gradients: Grad[Params], - params: Params, - state: State[Params, V] - )(using IsFloating[V]): (Params, State[Params, V]) = +class AdamW[V: IsFloating]( + val adam: Adam[V], + val weightDecayFactor: Tensor0[V] +) extends GradientOptimizer[V, [P] =>> AdamState[P, V]]: + + def init[P](params: P)(using TensorTree[P], FloatTree[P, V]): AdamState[P, V] = adam.init(params) + + def update[P]( + gradients: Grad[P], + params: P, + state: AdamState[P, V] + )(using TensorTree[P], FloatTree[P, V]): (P, AdamState[P, V]) = val α = adam.learningRate val `θₜ₋₁` = params val `λ'` = weightDecayFactor @@ -162,3 +168,33 @@ case class AdamW( val decayedParams = `θₜ₋₁` -- λ **! `θₜ₋₁` val (θₜ, adamState) = adam.update(gradients, decayedParams, state) (θₜ, adamState) + +case class LearningRateScheduleState[P, State[_]]( + step: Tensor0[Int32], + optState: State[P] +) +type LearningRateScheduleStateFor[State[_]] = [P] =>> LearningRateScheduleState[P, State] + +object LearningRateSchedule: + + def of[V: IsFloating, State[_]]( + vtype: VType[V] + )( + optF: Tensor0[V] => GradientOptimizer[V, State], + schedule: Tensor0[Int32] => Tensor0[V] + ): LearningRateSchedule[V, State] = + new LearningRateSchedule(optF, schedule) + +class LearningRateSchedule[V: IsFloating, State[_]](val optF: Tensor0[V] => GradientOptimizer[V, State], schedule: Tensor0[Int32] => Tensor0[V]) extends GradientOptimizer[V, LearningRateScheduleStateFor[State]]: + + def init[P](params: P)(using TensorTree[P], FloatTree[P, V]): LearningRateScheduleState[P, State] = + val step = Tensor0(1) + val opt = optF(schedule(step)) + LearningRateScheduleState(step, opt.init(params)) + + def update[P](gradients: Grad[P], params: P, state: LearningRateScheduleState[P, State])(using TensorTree[P], FloatTree[P, V]): (P, LearningRateScheduleState[P, State]) = + val step = state.step + val optState = state.optState + val opt = optF(schedule(step)) + val (newParams, newOptState) = opt.update(gradients, params, optState) + (newParams, LearningRateScheduleState(step + 1, newOptState)) diff --git a/core/src/main/scala/dimwit/tensor/TensorOps.scala b/core/src/main/scala/dimwit/tensor/TensorOps.scala index 3f7dc8d..ce1306c 100644 --- a/core/src/main/scala/dimwit/tensor/TensorOps.scala +++ b/core/src/main/scala/dimwit/tensor/TensorOps.scala @@ -28,9 +28,8 @@ object TensorOps: given [V](using ev1: IsFloating[V]): IsNumber[V] = ev1 given [V](using ev2: IsInteger[V]): IsNumber[V] = ev2 - @implicitNotFound("Operation only valid for Floating tensors.") - /** Type class marker for floating point types (Float32, Float64, etc.). */ + @implicitNotFound("Operation only valid for Floating tensors.") trait IsFloating[V] extends IsNumber[V], HasDType[V]: def dtype: DType diff --git a/core/src/main/scala/dimwit/tensortree/FloatTree.scala b/core/src/main/scala/dimwit/tensortree/FloatTree.scala index 4b3cd99..e68663d 100644 --- a/core/src/main/scala/dimwit/tensortree/FloatTree.scala +++ b/core/src/main/scala/dimwit/tensortree/FloatTree.scala @@ -142,3 +142,24 @@ object FloatTree: p.map([T <: Tuple] => (n: Labels[T]) ?=> (a: Tensor[T, V]) => a.asFloat(vtype)).asInstanceOf[F[NewV]] type FloatTreeFor[V] = [P] =>> FloatTree[P, V] + +/** A typeclass that proves P is a FloatTree, hiding the specific float type V + * from method signatures while keeping the evidence available. + */ +trait IsFloatTree[P]: + type V + given isFloating: IsFloating[V] + given floatTree: FloatTree[P, V] + given tensorTree: TensorTree[P] + +object IsFloatTree: + // The compiler automatically packages a FloatTree[P, V] into an IsFloatTree[P] + given pack[P, V0](using ft: FloatTree[P, V0], tt: TensorTree[P], isF: IsFloating[V0]): IsFloatTree[P] with + type V = V0 + val isFloating = isF + val floatTree = ft + val tensorTree = tt + + given unpackFloatTree[P](using isFT: IsFloatTree[P]): FloatTree[P, isFT.V] = isFT.floatTree + given unpackIsFloating[P](using isFT: IsFloatTree[P]): IsFloating[isFT.V] = isFT.isFloating + given unpackTensorTree[P](using isFT: IsFloatTree[P]): TensorTree[P] = isFT.tensorTree diff --git a/core/src/test/scala/dimwit/optimizer/GradientOptimizerSuite.scala b/core/src/test/scala/dimwit/optimizer/GradientOptimizerSuite.scala index dc8e2fb..02a87e4 100644 --- a/core/src/test/scala/dimwit/optimizer/GradientOptimizerSuite.scala +++ b/core/src/test/scala/dimwit/optimizer/GradientOptimizerSuite.scala @@ -10,18 +10,39 @@ class GradientOptimizerSuite extends DimwitTest: describe("GradientDescent"): it("should converge towards the minimum of f(x) = (x+1)^2 at x = -1"): - val optimizer = GradientDescent(learningRate = 0.1) + val optimizer = GradientDescent.of(VType[Float32])(learningRate = 0.1) val minX = optimizer.iterate(Tensor0(2.0f))(x => Grad(2 * (x + 1))).drop(1000).next() minX.item shouldBe -1.0f +- 0.1f + it("should compute the step (single step)"): + val optimizer = GradientDescent.of(VType[Float32])(learningRate = 0.1) + val initParams = Tensor0(2.0f) + val initState = optimizer.init(initParams) + + val grad = Grad(Tensor0(6.0f)) + val (nextParams, nextState) = optimizer.update(grad, initParams, initState) + + nextParams.item shouldBe 1.4f +- 1e-5f + + describe("LearningRateSchedule"): + it("basic"): + def learningRateSchedule(step: Tensor0[Int32]): Tensor0[Float32] = + if step.item <= 5 then Tensor0(0.0f) + else Tensor0(0.1f) + val optimizer = LearningRateSchedule.of(VType[Float32])(lr => GradientDescent(lr), learningRateSchedule) + val x1 = optimizer.iterate(Tensor0(2.0f))(x => Grad(2 * (x + 1))).drop(5).next() + x1.item shouldBe 2.0f +- 0.1 + val x2 = optimizer.iterate(Tensor0(2.0f))(x => Grad(2 * (x + 1))).drop(1000).next() + x2.item shouldBe -1.0f +- 0.1f + describe("Adam"): it("should converge towards the minimum of f(x) = (x+1)^2 at x = -1"): - val optimizer = Adam(learningRate = 0.1) + val optimizer = Adam.of(VType[Float32])(learningRate = 0.1f) val minX = optimizer.iterate(Tensor0(2.0f))(x => Grad(2 * (x + 1))).drop(1000).next() minX.item shouldBe -1.0f +- 0.1f it("should compute the exact momentum and velocity updates (single step)"): - val optimizer = Adam(learningRate = 0.1, b1 = 0.9, b2 = 0.999) + val optimizer = Adam.of(VType[Float32])(learningRate = 0.1f, b1 = 0.9f, b2 = 0.999f) val initParams = Tensor0(2.0f) val initState = optimizer.init(initParams) @@ -32,38 +53,39 @@ class GradientOptimizerSuite extends DimwitTest: nextState.momentums.item shouldBe 0.6f +- 1e-5f nextState.velocities.item shouldBe 0.036f +- 1e-5f nextState.b1.item shouldBe 0.9f +- 1e-5f - nextState.b2.item shouldBe 0.999f +- 1e-5f + nextState.b2.item shouldBe 0.999f +- 1e-5 describe("AdamW"): it("should converge towards the minimum of f(x) = (x+1)^2 at x = -1"): - val optimizer = AdamW(Adam(learningRate = 0.1), weightDecayFactor = 0.1) + val optimizer = AdamW(Adam.of(VType[Float32])(learningRate = 0.1f), weightDecayFactor = 0.1f) val minX = optimizer.iterate(Tensor0(2.0f))(x => Grad(2 * (x + 1))).drop(1000).next() minX.item shouldBe -1.0f +- 0.1f it("should apply decoupled weight decay (single step)"): - val adam = Adam(learningRate = 0.1) - val adamW = AdamW(adam, weightDecayFactor = 0.1) + val adam = Adam.of(VType[Float32])(learningRate = 0.1f) + val adamW = AdamW(adam, weightDecayFactor = 0.1f) - val initParams = Tensor0(2.0) - val grad = Grad(Tensor0(6.0)) + val initParams = Tensor0(2.0f) + val grad = Grad(Tensor0(6.0f)) val (adamParams, _) = adam.update(grad, initParams, adam.init(initParams)) val (adamWParams, _) = adamW.update(grad, initParams, adamW.init(initParams)) - adamWParams.item shouldBe (adamParams.item - 0.02) +- 1e-5 + adamWParams.item shouldBe (adamParams.item - 0.02f) +- 1e-5f describe("Lion"): it("should converge towards the minimum of f(x) = (x+1)^2"): - val optimizer = Lion(learningRate = 0.1) + val optimizer = Lion.of(VType[Float32])(learningRate = Tensor0(0.1f)) val minX = optimizer.iterate(Tensor0(2.0f))(x => Grad(2 * (x + 1))).drop(1000).next() minX.item shouldBe -1.0f +- 0.1f it("should compute the exact sign-based update and momentum (single step)"): - val optimizer = Lion(learningRate = 0.1, beta1 = 0.9, beta2 = 0.99) - val initParams = Tensor0(2.0) + val optimizer = Lion.of(VType[Float32])(learningRate = Tensor0(0.1f), beta1 = Tensor0(0.9f), beta2 = Tensor0(0.99f)) + val initParams = Tensor0(2.0f) val initMomentum = optimizer.init(initParams) - val grad = Grad(Tensor0(6.0)) - val (nextParams, nextMomentum) = optimizer.update(grad, initParams, initMomentum) + val grad = Grad(Tensor0(6.0f)) + val (nextParams, nextState) = optimizer.update(grad, initParams, initMomentum) - nextParams.item shouldBe 1.9d +- 1e-5d - nextMomentum.item shouldBe 0.06d +- 1e-5d + nextParams.item shouldBe 1.9f +- 1e-5f + nextState.momentums.item shouldBe 0.06f +- 1e-5f + nextState.step.item shouldBe 2 diff --git a/docs/quickstart.md b/docs/quickstart.md index 28569c8..fa4b058 100644 --- a/docs/quickstart.md +++ b/docs/quickstart.md @@ -11,6 +11,7 @@ Before we start exploring the features of DimWit, let's look at a simple example ```scala // main imports for basic tensor operations and automatic differentiation import dimwit.* +import dimwit.Conversions.given import dimwit.Autodiff.grad // TODO replace with cleaner import after PR is merged import dimwit.optimizer.GradientDescent // TODO replace with cleaner import after refactoring @@ -43,7 +44,7 @@ def fit(x: Tensor2[Batch, Feature, Float32], y: Tensor1[Batch, Float32]): Iterat val gradFn = grad(loss(x, y)) // gradient based optimization - val gd = GradientDescent(learningRate = 0.1f) // this is wrong, should be 0.1f not Tensor0 + val gd = GradientDescent.of(VType[Float32])(learningRate = 0.1f) // this is wrong, should be 0.1f not Tensor0 gd.iterate(p0)(gradFn) ``` @@ -219,10 +220,10 @@ tensor1 + tensor3 // Conflicting definitions: // val tensor1: // dimwit.tensor.Tensor[(MdocApp1.this.A, MdocApp1.this.B), -// dimwit.tensor.DType.Float32] in class MdocApp1 at line 63 and +// dimwit.tensor.DType.Float32] in class MdocApp1 at line 64 and // val tensor1: // dimwit.tensor.Tensor[(MdocApp1.this.A, MdocApp1.this.B), -// dimwit.tensor.DType.Float32] in class MdocApp1 at line 67 +// dimwit.tensor.DType.Float32] in class MdocApp1 at line 68 // // val tensor1 = Tensor(Shape(Axis[A] -> 3, Axis[B] -> 2)).fill(1.0f) // ^ diff --git a/examples/src/main/scala/dimwit/basic/LogisticRegression.scala b/examples/src/main/scala/dimwit/basic/LogisticRegression.scala index 239e24d..e7c7c88 100644 --- a/examples/src/main/scala/dimwit/basic/LogisticRegression.scala +++ b/examples/src/main/scala/dimwit/basic/LogisticRegression.scala @@ -125,7 +125,7 @@ object LogisticRegression: val trainLoss = jit(BinaryLogisticRegression.loss(trainingData, trainLabels)) val valLoss = jit(BinaryLogisticRegression.loss(valData, valLabels)) val learningRate = 5e-1f - val gd = GradientDescent(learningRate) + val gd = GradientDescent.of(VType[Float32])(learningRate) // Training loop val numiterations = 1000 diff --git a/examples/src/main/scala/dimwit/complex/VariationalAutoencoder.scala b/examples/src/main/scala/dimwit/complex/VariationalAutoencoder.scala index 73ca4a0..ea72c1f 100644 --- a/examples/src/main/scala/dimwit/complex/VariationalAutoencoder.scala +++ b/examples/src/main/scala/dimwit/complex/VariationalAutoencoder.scala @@ -209,34 +209,35 @@ object VariationalAutoencoderExample: losses.sum / batchSize.toFloat val batches = trainImages.chunk(Axis[TrainSample], numSamples / batchSize) - val optimizer = GradientDescent(learningRate = learningRate) - def trainBatch(trainKey: Random.Key, batch: Tensor3[TrainSample, Height, Width, Float32], params: Params): Params = + val optimizer = GradientDescent.of(VType[Float32])(learningRate = learningRate) + def trainBatch(trainKey: Random.Key, batch: Tensor3[TrainSample, Height, Width, Float32], params: Params, state: optimizer.State[Params]): (Params, optimizer.State[Params]) = val grads = grad(batchLoss(trainKey, batch))(params) - val (newParams, _) = optimizer.update(grads, params, ()) - newParams + val (newParams, newState) = optimizer.update(grads, params, state) + (newParams, newState) val (jitDonate, jitStep, jitReclaim) = jitDonating(trainBatch) - def trainEpoch(key: Random.Key, epoch: Int, params: Params): Params = + def trainEpoch(key: Random.Key, epoch: Int, params: Params, state: optimizer.State[Params]): (Params, optimizer.State[Params]) = val batchKeys = key.split(batches.size) jitReclaim( - batches.zip(batchKeys).foldLeft(jitDonate(params)): - case (batchParams, (batch, key)) => - jitStep(key, batch, batchParams) + batches.zip(batchKeys).foldLeft(jitDonate(params, state)): + case ((batchParams, state), (batch, key)) => + jitStep(key, batch, batchParams, state) ) val keysForEpochs = dataKey.split(numEpochs) val initialParams = Params(encoderParams, decoderParams).map([T <: Tuple] => (n: Labels[T]) ?=> (t: Tensor[T, Float32]) => t *! 0.1f) + val initState = optimizer.init(initialParams) - val trainedParams = (0 until numEpochs).foldLeft(initialParams): - case (params, epoch) => + val (trainedParams, _) = (0 until numEpochs).foldLeft(initialParams, initState): + case ((params, state), epoch) => timed(s"Evaluation $epoch/$numEpochs"): val lossValue = batchLoss(keysForEpochs(epoch), testImages)(params) println(s"Test loss in epoch $epoch: $lossValue") timed(s"Training $epoch/$numEpochs"): dimwit.gc() - trainEpoch(keysForEpochs(epoch), epoch, params) + trainEpoch(keysForEpochs(epoch), epoch, params, state) /* * Evaluation diff --git a/mdocs/AGENTS.md b/mdocs/AGENTS.md index 89adab9..e5e6496 100644 --- a/mdocs/AGENTS.md +++ b/mdocs/AGENTS.md @@ -677,6 +677,7 @@ println(s"Block Hessian shapes: ${h_x1x1.shape}, ${h_x1x2.shape}, ${h_x2x1.shape ```scala mdoc:reset:silent import dimwit.* +import dimwit.Conversions.given import dimwit.optimizer.{GradientDescent, GradientOptimizer} import dimwit.random.Random @@ -710,7 +711,7 @@ val lossFunc = mse(trainData, trainLabels) val gradFunc = Autodiff.grad(lossFunc) // Create optimizer -val optimizer = GradientDescent(learningRate = 0.01f) +val optimizer = GradientDescent.of(VType[Float32])(learningRate = 0.01f) // Training loop with iterator val trained = optimizer.iterate(initModelParams)(gradFunc) @@ -725,9 +726,10 @@ val trained = optimizer.iterate(initModelParams)(gradFunc) ```scala mdoc:silent import dimwit.optimizer.Lion +import dimwit.Conversions.given // enables implicit conversion from Float to Tensor[V] // Lion optimizer with momentum -val lionOptimizer = Lion(learningRate = 1e-3f, beta1 = 0.9f, beta2 = 0.99f, weightDecay = 0.0f) +val lionOptimizer = Lion.of(VType[Float32])(learningRate = 1e-3f, beta1 = 0.9f, beta2 = 0.99f, weightDecay = 0.0f) // Training with Lion val trainedLion = lionOptimizer.iterate(initModelParams)(gradFunc) @@ -739,6 +741,8 @@ val trainedLion = lionOptimizer.iterate(initModelParams)(gradFunc) ### Complete Training Example: Linear Regression ```scala mdoc:silent +import dimwit.Conversions.given // enables implicit conversion from Float to Tensor[V] + // Define problem dimensions trait Sample derives Label trait InputDim derives Label @@ -770,7 +774,7 @@ val initRegressionParams = RegressionParams(initSlope, initIntercept) // Train val regressionGrad = Autodiff.grad(regressionLoss(xData, yData)) -val gdOptimizer = GradientDescent(learningRate = 0.1f) +val gdOptimizer = GradientDescent.of(VType[Float32])(learningRate = 0.1f) val finalParams = gdOptimizer.iterate(initRegressionParams)(regressionGrad) .take(100) diff --git a/mdocs/docs/quickstart.md b/mdocs/docs/quickstart.md index 843046a..ea4c445 100644 --- a/mdocs/docs/quickstart.md +++ b/mdocs/docs/quickstart.md @@ -11,6 +11,7 @@ Before we start exploring the features of DimWit, let's look at a simple example ```scala mdoc:silent // main imports for basic tensor operations and automatic differentiation import dimwit.* +import dimwit.Conversions.given import dimwit.Autodiff.grad // TODO replace with cleaner import after PR is merged import dimwit.optimizer.GradientDescent // TODO replace with cleaner import after refactoring @@ -43,7 +44,7 @@ def fit(x: Tensor2[Batch, Feature, Float32], y: Tensor1[Batch, Float32]): Iterat val gradFn = grad(loss(x, y)) // gradient based optimization - val gd = GradientDescent(learningRate = 0.1f) // this is wrong, should be 0.1f not Tensor0 + val gd = GradientDescent.of(VType[Float32])(learningRate = 0.1f) // this is wrong, should be 0.1f not Tensor0 gd.iterate(p0)(gradFn) ```