graham/prop
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add --no-pretrain flag for direct QSO training comparison

Browse source
This commit is contained in:
Graham McIntire 2026-04-01 09:36:32 -05:00
parent 9537c97d1d
commit 489e188956
No known key found for this signature in database
GPG key ID: F4ABF488E6029E59
1 changed files with 44 additions and 30 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

74
lib/mix/tasks/propagation_train.ex
View file

@ -48,7 +48,8 @@ defmodule Mix.Tasks.PropagationTrain do
pretrain_epochs: :integer,
finetune_epochs: :integer,
pretrain_sample: :integer,
batch_size: :integer
batch_size: :integer,
no_pretrain: :boolean
]
)
@ -68,38 +69,53 @@ defmodule Mix.Tasks.PropagationTrain do
finetune_epochs = Keyword.get(opts, :finetune_epochs, 200)
pretrain_sample = Keyword.get(opts, :pretrain_sample, 500_000)
batch_size = Keyword.get(opts, :batch_size, 256)
skip_pretrain = Keyword.get(opts, :no_pretrain, false)
"=" |> String.duplicate(70) |> IO.puts()
IO.puts("PROPAGATION MODEL TRAINING (Two-Phase, 20 features)")
IO.puts("PROPAGATION MODEL TRAINING (20 features)")
"=" |> String.duplicate(70) |> IO.puts()
IO.puts("Phase 1: #{pretrain_epochs} epochs on #{pretrain_sample} algorithm scores")
IO.puts("Phase 2: #{finetune_epochs} epochs on real QSO data")
IO.puts("Started: #{DateTime.to_string(DateTime.utc_now())}\n")
# ── Phase 1: Pre-train on algorithm scores ──────────────────────
IO.puts("=" <> String.duplicate("─", 50))
IO.puts("PHASE 1: Pre-training on algorithm scores + solar + soundings")
IO.puts("=" <> String.duplicate("─", 50))
pretrained_state =
if skip_pretrain do
IO.puts("Skipping pre-training (--no-pretrain)")
nil
else
IO.puts("Phase 1: #{pretrain_epochs} epochs on #{pretrain_sample} algorithm scores")
IO.puts("Phase 2: #{finetune_epochs} epochs on real QSO data")
IO.puts("Started: #{DateTime.to_string(DateTime.utc_now())}\n")
{pt_features, pt_targets, pt_counts} = load_pretrain_data(pretrain_sample)
pt_n = elem(Nx.shape(pt_features), 0)
IO.puts("Loaded #{pt_n} samples:")
print_band_counts(pt_counts)
# ── Phase 1: Pre-train on algorithm scores ──────────────────
IO.puts("=" <> String.duplicate("─", 50))
IO.puts("PHASE 1: Pre-training on algorithm scores + solar + soundings")
IO.puts("=" <> String.duplicate("─", 50))
{pt_features, pt_targets} = shuffle(pt_features, pt_targets, pt_n)
{pt_train_x, pt_train_y, pt_val_x, pt_val_y, _, _} = split(pt_features, pt_targets, pt_n)
{pt_features, pt_targets, pt_counts} = load_pretrain_data(pretrain_sample)
pt_n = elem(Nx.shape(pt_features), 0)
IO.puts("Loaded #{pt_n} samples:")
print_band_counts(pt_counts)
IO.puts("\nPre-training for #{pretrain_epochs} epochs (lr=0.001)...")
{pt_features, pt_targets} = shuffle(pt_features, pt_targets, pt_n)
{pt_train_x, pt_train_y, pt_val_x, pt_val_y, _, _} = split(pt_features, pt_targets, pt_n)
{pretrained_state, pt_metrics} =
Model.train(pt_train_x, pt_train_y, epochs: pretrain_epochs, batch_size: batch_size)
IO.puts("\nPre-training for #{pretrain_epochs} epochs (lr=0.001)...")
print_loss(pt_metrics.final_loss)
print_eval("Pre-train val", Model.evaluate(pretrained_state, pt_val_x, pt_val_y))
{state, pt_metrics} =
Model.train(pt_train_x, pt_train_y, epochs: pretrain_epochs, batch_size: batch_size)
# ── Phase 2: Fine-tune on real QSO data ─────────────────────────
print_loss(pt_metrics.final_loss)
print_eval("Pre-train val", Model.evaluate(state, pt_val_x, pt_val_y))
state
end
# ── QSO training ────────────────────────────────────────────────
IO.puts("\n" <> "=" <> String.duplicate("─", 50))
IO.puts("PHASE 2: Fine-tuning on QSO + HRRR + solar + soundings")
if pretrained_state do
IO.puts("PHASE 2: Fine-tuning on QSO + HRRR + solar + soundings")
else
IO.puts("TRAINING: QSO + HRRR + solar + soundings (direct, lr=0.001)")
end
IO.puts("=" <> String.duplicate("─", 50))
{ft_features, ft_targets, ft_counts} = load_qso_data()
@ -110,15 +126,13 @@ defmodule Mix.Tasks.PropagationTrain do
{ft_features, ft_targets} = shuffle(ft_features, ft_targets, ft_n)
{ft_train_x, ft_train_y, ft_val_x, ft_val_y, ft_test_x, ft_test_y} = split(ft_features, ft_targets, ft_n)
IO.puts("\nFine-tuning for #{finetune_epochs} epochs (lr=0.0003)...")
lr = if pretrained_state, do: 0.0003, else: 0.001
IO.puts("\nTraining for #{finetune_epochs} epochs (lr=#{lr})...")
{trained_state, ft_metrics} =
Model.train(ft_train_x, ft_train_y,
epochs: finetune_epochs,
batch_size: batch_size,
learning_rate: 0.0003,
initial_state: pretrained_state
)
train_opts = [epochs: finetune_epochs, batch_size: batch_size, learning_rate: lr]
train_opts = if pretrained_state, do: [{:initial_state, pretrained_state} | train_opts], else: train_opts
{trained_state, ft_metrics} = Model.train(ft_train_x, ft_train_y, train_opts)
print_loss(ft_metrics.final_loss)
print_eval("Fine-tune val", Model.evaluate(trained_state, ft_val_x, ft_val_y))