swift_tensor_lr.swift

import TensorFlow

print(Tensor(3))

// y = - 3 * x + 7

let indices = -10...10

let xs = indices.map { x -> Float in Float(x) }
let ys = indices.map { x -> Float in Float(x) * (-3) + 7 + Float.random(in: -1.0...1.0) }

struct Params : ParameterAggregate {
	var w = Tensor<Float>(randomNormal: [1, 1])
	var b = Tensor<Float>(randomNormal: [1, 1])
}

var model = Params()

func error(_ model: Params, _ x: Tensor<Float>, _ y: Tensor<Float>) -> Tensor<Float> {
	let p = x * model.w + model.b
	let d = y - p
	return d * d
}

func predictAndGradient(_ model: Params, _ x: Tensor<Float>, _ y: Tensor<Float>) -> (Tensor<Float>, Params) {
	let (v, (grad, _, _)) = #valueAndGradient(error)(model, x, y)
	return (v, grad)
}

func train(_ model: inout Params) -> (Tensor<Float>, Params) {
	var av = Tensor<Float>(0)
	var ag = Params()
	ag.w = Tensor<Float>(0)
	ag.b = Tensor<Float>(0)

	for i in 0...20 {
		let (v, g) = predictAndGradient(model, Tensor(xs[i]), Tensor(ys[i]))
		av += v
		ag.w += g.w
		ag.b += g.b
	}

	model.w -= ag.w / 21 * 0.01
	model.b -= ag.b / 21 * 0.01
	
	return (av, model)
}
 
for i in 0...350 {
	print(train(&model))
}

print(5 * model.w + model.b) // should be around -8