4PL
Mix.install([
{:nx, "~> 0.8"},
{:exla, "~> 0.8"},
{:kino, "~> 0.14"},
{:kino_vega_lite, "~> 0.1"},
{:statistics, "~> 0.6"}
])
Nx.global_default_backend(EXLA.Backend)
データ作成
x_data = Nx.iota({20})
true_y_data =
Nx.divide(
0.5 - 7.3,
Nx.divide(x_data, 8)
|> Nx.pow(2.5)
|> Nx.add(1)
)
|> Nx.add(7.3)
y_data =
true_y_data
|> Nx.add(
1..20
|> Enum.map(fn _ -> Statistics.Distributions.Normal.rand() end)
|> Nx.tensor()
|> Nx.multiply(0.2)
)
plot_data =
%{
x: x_data |> Nx.to_flat_list(),
y: y_data |> Nx.to_flat_list()
}
true_plot_data =
%{
x: x_data |> Nx.to_flat_list(),
y: true_y_data |> Nx.to_flat_list()
}
VegaLite.new(width: 600)
|> VegaLite.layers([
VegaLite.new()
|> VegaLite.data_from_values(plot_data)
|> VegaLite.mark(:point)
|> VegaLite.encode_field(:x, "x", type: :quantitative)
|> VegaLite.encode_field(:y, "y", type: :quantitative),
VegaLite.new()
|> VegaLite.data_from_values(true_plot_data)
|> VegaLite.mark(:line, color: "#ff0000")
|> VegaLite.encode_field(:x, "x", type: :quantitative)
|> VegaLite.encode_field(:y, "y", type: :quantitative)
])
|> Kino.VegaLite.new()
トレーニング
defmodule FPL do
import Nx.Defn
defn pred({a, b, c, d}, x) do
(a - d) / (1.0 + Nx.pow(x / c, b)) + d
end
defn mse(yp, y) do
(yp - y)
|> Nx.pow(2)
|> Nx.mean()
end
defn loss(params, x, y) do
yp = pred(params, x)
mse(yp, y)
end
defn update({a, b, c, d} = params, x, y, lr) do
{grad_a, grad_b, grad_c, grad_d} = grad(params, &loss(&1, x, y))
{
a - grad_a * lr,
b - grad_b * lr,
c - grad_c * lr,
d - grad_d * lr
}
end
defn init_params do
{Nx.tensor(1.0), Nx.tensor(1.0), Nx.tensor(1.0), Nx.tensor(1.0)}
end
def loss_update({lvs, a, b, c, d}, x, y, lr) do
lv = FPL.loss({a, b, c, d}, x, y)
{a, b, c, d} = FPL.update({a, b, c, d}, x, y, lr)
{[Nx.to_number(lv) | lvs], a, b, c, d}
end
end
loss_widget =
VegaLite.new(width: 600)
|> VegaLite.mark(:line)
|> VegaLite.encode_field(:x, "x", type: :quantitative, title: "epoch")
|> VegaLite.encode_field(:y, "y", type: :quantitative, title: "loss")
|> Kino.VegaLite.new()
fpl_widget =
VegaLite.new(width: 600)
|> VegaLite.layers([
VegaLite.new()
|> VegaLite.data_from_values(plot_data)
|> VegaLite.mark(:point)
|> VegaLite.encode_field(:x, "x", type: :quantitative)
|> VegaLite.encode_field(:y, "y", type: :quantitative),
VegaLite.new()
|> VegaLite.data_from_values(true_plot_data)
|> VegaLite.mark(:line, color: "#ff0000")
|> VegaLite.encode_field(:x, "x", type: :quantitative)
|> VegaLite.encode_field(:y, "y", type: :quantitative),
VegaLite.new()
|> VegaLite.mark(:line)
|> VegaLite.encode_field(:x, "x", type: :quantitative)
|> VegaLite.encode_field(:y, "y", type: :quantitative)
])
|> Kino.VegaLite.new()
Kino.VegaLite.clear(loss_widget)
Kino.VegaLite.clear(fpl_widget)
Kino.Layout.grid([loss_widget, fpl_widget], columns: 1)
update_plots = fn {epoch, lvs, a, b, c, d} ->
loss_plot_data =
1..epoch
|> Enum.zip(Enum.reverse(lvs))
|> Enum.map(fn {x, y} -> %{x: x, y: y} end)
Kino.VegaLite.clear(loss_widget)
Kino.VegaLite.push_many(loss_widget, loss_plot_data)
yl_data = FPL.pred({a, b, c, d}, x_data)
fpl_plot_data =
Enum.zip(
x_data |> Nx.to_flat_list(),
yl_data |> Nx.to_flat_list()
)
|> Enum.map(fn {x, y} -> %{x: x, y: y} end)
Kino.VegaLite.clear(fpl_widget)
Kino.VegaLite.push_many(fpl_widget, fpl_plot_data)
end
{a, b, c, d} = FPL.init_params()
epochs = 2500
lr = 0.02
Enum.reduce(1..epochs, {[], a, b, c, d}, fn epoch, acc ->
{lvs, a, b, c, d} = FPL.loss_update(acc, x_data, y_data, lr)
if rem(epoch, 10) == 0 do
update_plots.({epoch, lvs, a, b, c, d})
end
{lvs, a, b, c, d}
end)