Skip to content

clean up, complete and comment logistic regression example #43

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
marcelluethi merged 1 commit into from
Jan 19, 2026
Merged

clean up, complete and comment logistic regression example #43

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
121 changes: 78 additions & 43 deletions examples/src/main/scala/basic/LogisticRegression.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,95 +1,130 @@
package examples.basic

import dimwit.*
import nn.*
import nn.ActivationFunctions.{relu, sigmoid}
import dimwit.Conversions.given
import nn.*
import nn.ActivationFunctions.{sigmoid, relu}
import dimwit.random.Random
import dimwit.stats.Normal

object LogisticRegression:

// Define labels for tensor axes
trait Sample derives Label
trait Feature derives Label

object BinaryLogisticRegression:
case class Params(
linearMap: LinearMap.Params[Feature]
)

// Define a binary logistic regression model
case class BinaryLogisticRegression(
params: BinaryLogisticRegression.Params
) extends Function[Tensor1[Feature, Float], Tensor0[Boolean]]:
private val linear = LinearMap(params.linearMap)
def logits(input: Tensor1[Feature, Float]): Tensor0[Float] = linear(input)
def probits(input: Tensor1[Feature, Float]): Tensor0[Float] = sigmoid(logits(input))
def apply(input: Tensor1[Feature, Float]): Tensor0[Boolean] = logits(input) >= Tensor0(0f)

def logits(input: Tensor1[Feature, Float]): Tensor0[Float] =
params.weights.dot(Axis[Feature])(input) + params.bias

def probits(input: Tensor1[Feature, Float]): Tensor0[Float] =
sigmoid(logits(input))

def apply(input: Tensor1[Feature, Float]): Tensor0[Boolean] =
logits(input) >= Tensor0(0f)

// Parameters are, by convention, defined in the companion object
object BinaryLogisticRegression:
case class Params(
weights: Tensor1[Feature, Float],
bias: Tensor0[Float]
) derives ToPyTree, FloatTensorTree

// The loss is a simple binary cross-entropy loss
def loss(data: Tensor2[Sample, Feature, Float], labels: Tensor1[Sample, Boolean])(params: BinaryLogisticRegression.Params)
: Tensor0[Float] =

// Create the model with the given parameters
val model = BinaryLogisticRegression(params)

// Compute the logistic loss for the model over the dataset
val losses = zipvmap(Axis[Sample])(data, labels.asFloat):
case (sample, label) =>
val logits = model.logits(sample)
relu(logits) - logits * label + ((-logits.abs).exp + 1f).log
losses.mean

def main(args: Array[String]): Unit =

// we need two keys. One for initializing parameters,
// the other for shuffling data
val (initKey, shuffleKey) = Random.Key(42).split2()

// Load and preprocess the penguin dataset
val df = PenguinCSV
.parse("./data/penguins.csv")
.filter(row => row.species != 2)

val dfShuffled = scala.util.Random.shuffle(df)

val featureData = dfShuffled.map { row =>
val featureData = df.map { row =>
Array(
row.flipper_length_mm.toFloat,
row.bill_length_mm.toFloat,
row.bill_depth_mm.toFloat,
row.body_mass_g.toFloat
)
}.toArray
val labelData = dfShuffled.map(_.species).toArray.map {

val labelData = df.map(_.species).toArray.map {
case 1 => true
case 0 => false
}

val dataUnnormalized = Tensor2(Axis[Sample], Axis[Feature]).fromArray(featureData)
val dataLabels = Tensor1(Axis[Sample]).fromArray(labelData)
val numSamples = featureData.length
val numFeatures = featureData.head.length

// Convert the data into tensors
val dataInitial = Tensor2(Axis[Sample], Axis[Feature]).fromArray(featureData)
val labelsInitial = Tensor1(Axis[Sample]).fromArray(labelData)

// Create a permutation to shuffle data and
// split the permutation indices into training and validation
val perm = Random.permutation(dataInitial.shape.dim(Axis[Sample]))(shuffleKey)
val testSplitRatio = 0.4f
val splitIndex = (numSamples * testSplitRatio).toInt
val trainPerm = perm.slice(Axis[Sample] -> (0 until splitIndex))
val testPerm = perm.slice(Axis[Sample] -> (splitIndex until numSamples))

// TODO implement split
val (trainingDataUnnormalized, valDataUnnormalized) = (dataUnnormalized, dataUnnormalized)
val (trainLabels, valLabels) = (dataLabels, dataLabels)
// Use the permutations to get our training and validation data
val trainingDataUnnormalized = dataInitial.take(Axis[Sample])(trainPerm)
val valDataUnnormalized = dataInitial.take(Axis[Sample])(testPerm)
val trainLabels = labelsInitial.take(Axis[Sample])(trainPerm)
val valLabels = labelsInitial.take(Axis[Sample])(testPerm)

def calcMeanAndStd(t: Tensor2[Sample, Feature, Float]): (Tensor1[Feature, Float], Tensor1[Feature, Float]) =
val mean = t.vmap(Axis[Feature])(_.mean)
val std = zipvmap(Axis[Feature])(t, mean):
case (x, m) =>
val epsilon = 1e-6f
(x -! m).pow(2f).mean.sqrt + epsilon
// x.vmap(Axis[Sample])(xi => (xi - m).pow(2)).mean.sqrt + epsilon
(mean, std)

def standardizeData(mean: Tensor1[Feature, Float], std: Tensor1[Feature, Float])(data: Tensor2[Sample, Feature, Float]): Tensor2[Sample, Feature, Float] =
data.vapply(Axis[Feature])(feature => (feature - mean) / std)
// (data :- mean) :/ std
def standardizeData(mean: Tensor1[Feature, Float], std: Tensor1[Feature, Float])(data: Tensor2[Sample, Feature, Float])
: Tensor2[Sample, Feature, Float] =
(data -! mean) /! std

// Standardize the training and validation data
val (trainMean, trainStd) = calcMeanAndStd(trainingDataUnnormalized)
val trainingData = standardizeData(trainMean, trainStd)(trainingDataUnnormalized)
val valData = standardizeData(trainMean, trainStd)(valDataUnnormalized)

val (dataKey, trainKey) = Random.Key(42).split2()
val (initKey, restKey) = trainKey.split2()
val (lossKey, sampleKey) = restKey.split2()

def loss(data: Tensor2[Sample, Feature, Float])(params: BinaryLogisticRegression.Params): Tensor0[Float] =
val model = BinaryLogisticRegression(params)
val losses = zipvmap(Axis[Sample])(data, trainLabels.asFloat):
case (sample, label) =>
val logits = model.logits(sample)
relu(logits) - logits * label + ((-logits.abs).exp + 1f).log
losses.mean

// Initialize model parameters
val initParams = BinaryLogisticRegression.Params(
LinearMap.Params(initKey)(dataUnnormalized.dim(Axis[Feature]))
weights = Normal.standardNormal(Shape(Axis[Feature] -> numFeatures)).sample(initKey) *! 0.01f,
bias = Tensor0(0f)
)

val trainLoss = loss(trainingData)
val valLoss = loss(valData)
val learningRate = 3e-1f
// Setting up the loss functions and optimizer.
// We always jit the loss functions for performance.
val trainLoss = jit(BinaryLogisticRegression.loss(trainingData, trainLabels))
val valLoss = jit(BinaryLogisticRegression.loss(valData, valLabels))
val learningRate = 5e-1f
val gd = GradientDescent(learningRate)

// Training loop
val numiterations = 1000
val trainTrajectory = gd.iterate(initParams)(Autodiff.grad(trainLoss))
val finalParams = trainTrajectory.zipWithIndex
.tapEach:
Expand All @@ -105,8 +140,8 @@ object LogisticRegression:
).mkString(", ")
)
.map((params, _) => params)
.drop(2500)
.next()
.drop(numiterations - 1) // we are only interested in the final parameters
.next() // consume the value we are actually interested in

val finalModel = BinaryLogisticRegression(finalParams)
val predictions = trainingData.vmap(Axis[Sample])(finalModel.probits)
Expand Down