"""Classes and functions for training GradNet magnetic models."""
import random
from collections.abc import Callable
from pathlib import Path
from typing import Literal, cast
import numpy as np
import torch
import torch.nn.functional as F
from motulator.drive.utils import BaseValues
from torch import Tensor
from torch.utils.data import DataLoader
from tqdm import trange
import motulator_gradnet._gn as gn
from motulator_gradnet._dataset import BaseDataset, SpatialHarmonicsDataset, get_loader
# %%
class Trainer:
"""
Trainer class for GradNet models.
Parameters
----------
model : GradNet
The GradNet model to be trained.
data_loader : DataLoader
DataLoader providing the training data.
dataset : BaseDataset | SpatialHarmonicsDataset
Dataset used for training.
mode : {"current_map", "flux_map"}
Which map to learn.
k : int | None
Spatial harmonics order. If None, no harmonics are used.
lr : float
Learning rate for the optimizer.
i_low : float, optional
Current limit (p.u.) below which the data points are weighted by `weight_low`.
weight_low : float, optional
Loss weight of the data points with current magnitude below `i_low`. The
default 1.0 gives the unweighted loss.
"""
def __init__(
self,
model: gn.GradNet,
data_loader: DataLoader,
dataset: BaseDataset | SpatialHarmonicsDataset,
mode: Literal["current_map", "flux_map"],
k: int | None,
lr: float,
i_low: float = 0.0,
weight_low: float = 1.0,
) -> None:
self.model = model
self.data_loader = data_loader
self.dataset = dataset
self.mode = mode
self.k = k
self.i_low = i_low
self.weight_low = weight_low
# Mean weight over the whole training set, which normalizes the loss so that
# each batch gives an unbiased estimate of the weighted loss of the whole set
self.w_mean = 1.0
if weight_low != 1.0:
w = self._point_weights(dataset.i_d, dataset.i_q)
self.w_mean = float(w.mean())
self.optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
def train_epoch(self) -> float:
"""Train the model for one epoch and return total loss."""
self.model.train()
total_loss = 0.0
for batch in self.data_loader:
self.optimizer.zero_grad()
loss = self._compute_batch_loss(batch)
loss.backward()
self.optimizer.step()
total_loss += loss.item()
return total_loss
def _compute_batch_loss(self, batch):
# Move batch to the same device as the model
batch = tuple(x.to(next(self.model.parameters()).device) for x in batch)
if self.k is None:
return self._compute_loss_no_harmonics(batch)
return self._compute_loss_with_harmonics(batch)
def _point_weights(self, i_d: Tensor, i_q: Tensor) -> Tensor:
"""Return the loss weights of the data points with the given currents."""
is_low = i_d**2 + i_q**2 < self.i_low**2
return torch.where(is_low, self.weight_low, 1.0)
def _weights(self, batch: tuple[Tensor, ...]) -> Tensor | None:
"""Return the normalized loss weights of the batch, or None if unweighted."""
if self.weight_low == 1.0:
return None
# The current of the data point is used in both modes
return self._point_weights(batch[2], batch[3]) / self.w_mean
def _compute_loss_no_harmonics(self, batch: tuple[Tensor, ...]) -> Tensor:
"""Compute loss for a batch without spatial harmonics."""
mode = cast(Literal["current_map", "flux_map"], self.mode)
inputs, targets = self.dataset.prepare_batch(batch, mode)
output = self.model(inputs)
return _mse_loss(output, targets, self._weights(batch))
def _compute_loss_with_harmonics(self, batch: tuple[Tensor, ...]) -> Tensor:
"""Compute loss for a batch with spatial harmonics."""
mode = cast(Literal["current_map", "flux_map"], self.mode)
k = cast(int, self.k)
inputs, targets = self.dataset.prepare_batch(batch, mode, k)
_, _, _, _, theta_m, tau_m_theta = batch
cos_t = torch.cos(k * theta_m)
sin_t = torch.sin(k * theta_m)
output = self.model(inputs)
dW_dcos = output[:, 2]
dW_dsin = output[:, 3]
dW_dtheta = k * (cos_t * dW_dsin - sin_t * dW_dcos)
# Dataset is normalized to per-unit values
if mode == "current_map":
tau_m_theta_pred = -dW_dtheta
else:
tau_m_theta_pred = dW_dtheta
w = self._weights(batch)
loss_main = _mse_loss(output[:, :2], targets, w)
loss_tau = _mse_loss(tau_m_theta_pred, tau_m_theta, w)
return loss_main + loss_tau
def _mse_loss(output: Tensor, targets: Tensor, w: Tensor | None) -> Tensor:
"""Mean squared error, weighted over the data points if `w` is given."""
if w is None:
return F.mse_loss(output, targets)
err2 = (output - targets) ** 2
if err2.ndim > 1:
err2 = err2.mean(dim=1)
return (w * err2).mean()
# %%
[docs]
def train_gradnet( # noqa: PLR0913
dataset_path: str | Path,
base: BaseValues,
is_flux_map=False,
k: int | None = None,
num_modules: int = 1,
embed_dim: int = 12,
batch_size: int = 128,
epochs: int = 2000,
lr: float = 1e-3,
save_model_path: str | Path | None = None,
subsample: int = 1,
activation: Callable[[], torch.nn.Module] | None = None,
device: torch.device | None = None,
*,
i_low: float = 0.0,
weight_low: float = 1.0,
) -> None:
"""
Train and save the GradNet model.
Parameters
----------
dataset_path : str | Path
Path to the training data file (npz format).
base : ~motulator.drive.utils.BaseValues
Base values of the target machine. The training data are normalized with the
base current and flux linkage, which are also stored in the model.
is_flux_map : bool, optional
Whether the model is a flux map or current map, defaults to False.
k : int | None, optional
Spatial harmonics order. If None, no harmonics are used, defaults to None.
num_modules : int, optional
Number of GradNet modules, defaults to 1.
embed_dim : int, optional
Embedding dimension for the GradNet modules, defaults to 12.
batch_size : int, optional
Batch size for training, defaults to 128.
epochs : int, optional
Number of training epochs, defaults to 2000.
lr : float, optional
Learning rate for the optimizer, defaults to 1e-3.
save_model_path : str | Path | None, optional
Path to save the trained model as an npz file (see `to_function`). If None,
saves to `model.npz` in current directory.
subsample : int, optional
Subsampling factor for the training data, defaults to 1 (uses all data).
activation : Callable[[], torch.nn.Module] | None, optional
Activation function factory, defaults to PNormGradient. The trained model is
saved as a numpy function, which supports the activations in `ACTIVATIONS`.
device : torch.device | None, optional
Device to use for training. If None, automatically selects CUDA if available,
otherwise CPU.
i_low : float, optional
Current limit (A, peak value) defining the low-current region, defaults to 0.
weight_low : float, optional
Loss weight of the data points whose current magnitude is below `i_low`,
defaults to 1 (unweighted loss). The low-current region is determined by the
current of the data point for both flux and current maps. A weight above 1
improves the accuracy at low currents, where a uniform data grid gives only a
few points. The weights are normalized with their mean over the training set,
so the loss of each batch is an unbiased estimate of the weighted loss of the
whole set, which corresponds to repeating those points `weight_low` times in the
dataset, except that the number of steps per epoch is unchanged.
"""
# Set random seed for reproducibility
seed = 42
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(seed)
if device is None:
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
# Initialize model and data
mode = "flux_map" if is_flux_map else "current_map"
if activation is None:
activation = gn.PNormGradient
in_dim = 4 if k is not None else 2
mu_dim = 2
# Per-unit data, normalized with the given base values
data_loader = get_loader(
dataset_path=dataset_path,
base=base,
k=k,
batch_size=batch_size,
subsample=subsample,
)
dataset = cast(BaseDataset | SpatialHarmonicsDataset, data_loader.dataset)
model = gn.GradNet(
in_dim,
mu_dim,
num_modules,
embed_dim,
activation,
mu_log0=-1.0,
psi_base=dataset.psi_base,
i_base=dataset.i_base,
).to(device)
gn.to_function(model) # Raises if the model cannot be saved
_run_training_loop(
model,
data_loader,
mode,
k,
lr,
epochs,
activation,
description="Training",
i_low=i_low / dataset.i_base,
weight_low=weight_low,
)
# Print model parameters and weights
print("\nModel parameters and weights:\n" + "-" * 40)
for name, param in model.named_parameters():
print(f"{name}: {param.shape}\n{param.data}\n{'-' * 40}")
# Save the model as a numpy function
if save_model_path is None:
save_model_path = Path.cwd() / "model.npz"
save_model_path = Path(save_model_path)
save_model_path.parent.mkdir(parents=True, exist_ok=True)
gn.to_function(model).save(save_model_path)
def _run_training_loop(
model,
data_loader,
mode,
k,
lr,
epochs,
activation,
description="Epoch",
i_low=0.0,
weight_low=1.0,
) -> None:
"""Helper to run training loop for a given data loader."""
dataset = cast(BaseDataset | SpatialHarmonicsDataset, data_loader.dataset)
_print_training_info(model, dataset, activation)
trainer = Trainer(model, data_loader, dataset, mode, k, lr, i_low, weight_low)
for epoch in trange(epochs, desc=description):
total_loss = trainer.train_epoch()
if epoch % (epochs // 10 if epochs > 10 else 1) == 0 or epoch == epochs - 1:
avg_loss = total_loss / len(data_loader)
print(f"[{description} {epoch}] Loss: {avg_loss:.6f}")
# %%
def _print_training_info(
model: gn.GradNet,
dataset: BaseDataset | SpatialHarmonicsDataset,
activation: Callable[[], torch.nn.Module],
) -> None:
"""Print information about the training configuration."""
print(f"Activation: {activation.__name__}")
n_trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"Trainable parameters: {n_trainable}")
print(
f"Normalization factors: "
f"i_base={dataset.i_base:.2f}, "
f"psi_base={dataset.psi_base:.2f}, "
f"tau_base={1.5 * dataset.psi_base * dataset.i_base:.2f}"
)