diff --git a/examples/src/main/scala/basic/LogisticRegression.scala b/examples/src/main/scala/basic/LogisticRegression.scala index b7cf5ac5..fe7dd2ad 100644 --- a/examples/src/main/scala/basic/LogisticRegression.scala +++ b/examples/src/main/scala/basic/LogisticRegression.scala @@ -1,38 +1,65 @@ 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, @@ -40,17 +67,32 @@ object LogisticRegression: 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) @@ -58,38 +100,31 @@ object LogisticRegression: 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: @@ -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)