64 lines
1.5 KiB
Nim
64 lines
1.5 KiB
Nim
import arraymancer
|
|
import std/[strformat, times]
|
|
|
|
# XOR dataset — [4,2] inputs, [4,1] targets
|
|
let xData = [[0.0'f32, 0.0'f32],
|
|
[0.0'f32, 1.0'f32],
|
|
[1.0'f32, 0.0'f32],
|
|
[1.0'f32, 1.0'f32]].toTensor()
|
|
|
|
let yData = [[0.0'f32], [1.0'f32], [1.0'f32], [0.0'f32]].toTensor()
|
|
|
|
var ctx = newContext Tensor[float32]
|
|
|
|
let l1 = ctx.init(Linear[float32], 2, 4)
|
|
let l2 = ctx.init(Linear[float32], 4, 1)
|
|
|
|
type XorNet = object
|
|
l1: Linear[float32]
|
|
l2: Linear[float32]
|
|
|
|
var net = XorNet(l1: l1, l2: l2)
|
|
var adam = optimizer(net, Adam, learning_rate = 0.01'f32)
|
|
|
|
let t0 = cpuTime()
|
|
|
|
for epoch in 1..1000:
|
|
let xVar = ctx.variable(xData)
|
|
let yVar = ctx.variable(yData)
|
|
|
|
# Forward
|
|
let h = l1.forward(xVar).sigmoid()
|
|
let pred = l2.forward(h).sigmoid()
|
|
|
|
# MSE loss: mean((pred - y)^2)
|
|
let diff = pred - yVar
|
|
let loss = mean(diff *. diff)
|
|
|
|
if epoch mod 200 == 0:
|
|
echo &"Epoch {epoch:4d} loss = {loss.value[0]:.6f}"
|
|
|
|
loss.backprop()
|
|
adam.update()
|
|
|
|
let trainTime = cpuTime() - t0
|
|
|
|
# Inference (no grad)
|
|
echo ""
|
|
echo "Final predictions:"
|
|
let xVar = ctx.variable(xData)
|
|
let hInf = l1.forward(xVar).sigmoid()
|
|
let preds = l2.forward(hInf).sigmoid()
|
|
|
|
var allCorrect = true
|
|
for i in 0..<4:
|
|
let p = preds.value[i, 0]
|
|
let rounded = if p >= 0.5'f32: 1 else: 0
|
|
let expected = [0,1,1,0][i]
|
|
if rounded != expected: allCorrect = false
|
|
echo &" [{xData[i,0].int},{xData[i,1].int}] -> {p:.4f} (rounded: {rounded}, expected: {expected})"
|
|
|
|
echo ""
|
|
echo &"All correct: {allCorrect}"
|
|
echo &"Training time: {trainTime:.2f}s"
|