Add model training options and show model predictions on plot

hyperparameters.py

@@ -0,0 +1,161 @@
+from dataclasses import dataclass, fields
+import gradio as gr
+
+
+@dataclass(frozen=True)
+class SgdHyperparameters:
+    learning_rate: float = 0.01
+    momentum: float = 0.0
+    weight_decay: float = 0.0
+    batch_size: int = 32
+
+
+@dataclass(frozen=True)
+class AdamHyperparameters:
+    learning_rate: float = 0.001
+    beta1: float = 0.9
+    beta2: float = 0.999
+    weight_decay: float = 0.0
+    batch_size: int = 32
+
+
+class Hyperparameters:
+    def __init__(
+        self,
+        optimizer: str = "SGD",
+        sgd_params: SgdHyperparameters = SgdHyperparameters(),
+        adam_params: AdamHyperparameters = AdamHyperparameters(),
+    ):
+        self.optimizer = optimizer
+        self.sgd_params = sgd_params
+        self.adam_params = adam_params
+
+    def update(self, **kwargs):
+        return Hyperparameters(
+            optimizer=kwargs.get("optimizer", self.optimizer),
+            sgd_params=kwargs.get("sgd_params", self.sgd_params),
+            adam_params=kwargs.get("adam_params", self.adam_params),
+        )
+
+    def __hash__(self):
+        return hash((self.optimizer, self.sgd_params, self.adam_params))
+
+    @property
+    def batch_size(self):
+        if self.optimizer == "SGD":
+            return self.sgd_params.batch_size
+        elif self.optimizer == "Adam":
+            return self.adam_params.batch_size
+        else:
+            raise ValueError(f"Unknown optimizer: {self.optimizer}")
+
+
+class HyperparametersView:
+    def update_optimizer_type(self, state: Hyperparameters, optimizer: str):
+        state = state.update(optimizer=optimizer)
+        return (
+            state,
+            gr.update(visible=(optimizer == "SGD")),
+            gr.update(visible=(optimizer == "Adam")),
+        )
+
+    def update_sgd_hyperparameters(
+        self,
+        state: Hyperparameters,
+        sgd_learning_rate: float,
+        sgd_momentum: float,
+        sgd_weight_decay: float,
+        sgd_batch_size: int,
+    ):
+        sgd_params = SgdHyperparameters(
+            learning_rate=sgd_learning_rate,
+            momentum=sgd_momentum,
+            weight_decay=sgd_weight_decay,
+            batch_size=sgd_batch_size,
+        )
+        state = state.update(sgd_params=sgd_params)
+        return state
+
+    def update_adam_hyperparameters(
+        self,
+        state: Hyperparameters,
+        adam_learning_rate: float,
+        adam_beta1: float,
+        adam_beta2: float,
+        adam_weight_decay: float,
+        adam_batch_size: int,
+    ):
+        adam_params = AdamHyperparameters(
+            learning_rate=adam_learning_rate,
+            beta1=adam_beta1,
+            beta2=adam_beta2,
+            weight_decay=adam_weight_decay,
+            batch_size=adam_batch_size,
+        )
+        state = state.update(adam_params=adam_params)
+        return state
+
+    def build(self, state: gr.State):
+        hyper = state.value
+        with gr.Column():
+            optimizer_select = gr.Dropdown(
+                choices=["SGD", "Adam"],
+                value=hyper.optimizer,
+                label="Optimizer",
+                interactive=True,
+            )
+
+            with gr.Group(visible=(hyper.optimizer == "SGD")) as sgd_box:
+                sgd_components = {}
+                with gr.Row():
+                    for f in fields(hyper.sgd_params):
+                        sgd_components[f.name] = gr.Number(
+                            value=getattr(hyper.sgd_params, f.name),
+                            label=f.name.replace("_", " ").title(),
+                            interactive=True,
+                        )
+
+            with gr.Group(visible=(hyper.optimizer == "Adam")) as adam_box:
+                adam_components = {}
+                with gr.Row():
+                    for f in fields(hyper.adam_params):
+                        adam_components[f.name] = gr.Number(
+                            value=getattr(hyper.adam_params, f.name),
+                            label=f.name.replace("_", " ").title(),
+                            interactive=True,
+                        )
+
+        optimizer_select.change(
+            fn=self.update_optimizer_type,
+            inputs=[state, optimizer_select],
+            outputs=[state, sgd_box, adam_box],
+        )
+
+        for name, component in sgd_components.items():
+            component.submit(
+                fn=self.update_sgd_hyperparameters,
+                inputs=[
+                    state,
+                    sgd_components["learning_rate"],
+                    sgd_components["momentum"],
+                    sgd_components["weight_decay"],
+                    sgd_components["batch_size"],
+                ],
+                outputs=[state],
+            )
+
+        for name, component in adam_components.items():
+            component.submit(
+                fn=self.update_adam_hyperparameters,
+                inputs=[
+                    state,
+                    adam_components["learning_rate"],
+                    adam_components["beta1"],
+                    adam_components["beta2"],
+                    adam_components["weight_decay"],
+                    adam_components["batch_size"],
+                ],
+                outputs=[state],
+            )
+
+

mlp_visualizer.py

@@ -1,5 +1,5 @@
 from collections import deque
-from dataclasses import dataclass, replace
+from dataclasses import dataclass, fields
 import functools
 from pathlib import Path
 import pickle
@@ -31,6 +31,68 @@ logger = logging.getLogger("ELVIS")
 
 from architecture import Architecture, ArchitectureView
 from dataset import Dataset, DatasetView, get_function
+from hyperparameters import Hyperparameters, HyperparametersView
+
+
+@dataclass
+class TrainState:
+    model: nn.Module
+    optimizer: torch.optim.Optimizer
+
+
+def init_model(architecture: Architecture) -> nn.Module:
+    input_size = 1
+    output_size = 1
+
+    layers = []
+    for hidden_units, activation in zip(architecture.hidden_units, architecture.activations):
+        layers.append(nn.Linear(input_size, hidden_units))
+
+        if activation == "ReLU":
+            layers.append(nn.ReLU())
+        elif activation == "Sigmoid":
+            layers.append(nn.Sigmoid())
+        elif activation == "Tanh":
+            layers.append(nn.Tanh())
+        elif activation == "LeakyReLU":
+            layers.append(nn.LeakyReLU())
+        elif activation == "ELU":
+            layers.append(nn.ELU())
+        elif activation == "GELU":
+            layers.append(nn.GELU())
+        elif activation == "Identity":
+            layers.append(nn.Identity())
+        else:
+            raise ValueError(f"Unknown activation: {activation}")
+
+        input_size = hidden_units
+
+    layers.append(nn.Linear(input_size, output_size))
+    model = nn.Sequential(*layers)
+    return model
+
+
+def init_optimizer(
+    model: nn.Module,
+    hyperparameters: Hyperparameters,
+) -> torch.optim.Optimizer:
+    if hyperparameters.optimizer == "SGD":
+        opt = torch.optim.SGD(
+            model.parameters(),
+            lr=hyperparameters.sgd_params.learning_rate,
+            momentum=hyperparameters.sgd_params.momentum,
+            weight_decay=hyperparameters.sgd_params.weight_decay,
+        )
+    elif hyperparameters.optimizer == "Adam":
+        opt = torch.optim.Adam(
+            model.parameters(),
+            lr=hyperparameters.adam_params.learning_rate,
+            betas=(hyperparameters.adam_params.beta1, hyperparameters.adam_params.beta2),
+            weight_decay=hyperparameters.adam_params.weight_decay,
+        )
+    else:
+        raise ValueError(f"Unknown optimizer: {hyperparameters.optimizer}")
+    return opt
 
 
 class MlpVisualizer:
@@ -45,7 +107,7 @@ class MlpVisualizer:
     display: none;
 }"""
 
-    def plot(self, dataset: Dataset, architecture: Architecture) -> Image.Image:
+    def plot(self, dataset: Dataset, train_state: TrainState) -> Image.Image:
         print("Plotting")
         t1 = time.time()
         fig = plt.figure(figsize=(self.canvas_width / 100., self.canvas_height / 100.0), dpi=100)
@@ -57,7 +119,7 @@ class MlpVisualizer:
         if dataset.mode == "generate":
             x_test, y_test = get_function(dataset.function, xlim=(-2, 2), nsample=100)
 
-        # y_pred = self.model(torch.from_numpy(x_test).float()).detach().numpy()
+        y_pred = train_state.model(torch.from_numpy(x_test).float()).detach().numpy()
 
         # plot
         fig, ax = plt.subplots(figsize=(8, 8))
@@ -76,7 +138,7 @@ class MlpVisualizer:
         if dataset.mode == "generate":
             plt.plot(x_test.flatten(), y_test, label='true function', color=self.plot_cmap(1))
 
-        if False:
+        if True:
             plt.plot(x_test.flatten(), y_pred, linestyle="--", label='prediction', color=self.plot_cmap(2))
 
         plt.legend()
@@ -92,6 +154,77 @@ class MlpVisualizer:
 
         return img
 
+    def update_dataset(
+        self,
+        dataset: Dataset,
+        architecture: Architecture,
+        hyperparameters: Hyperparameters,
+    ):
+        print("Updating dataset")
+        new_model = init_model(architecture)
+        new_optimizer = init_optimizer(new_model, hyperparameters)
+        new_train_state = TrainState(new_model, new_optimizer)
+        new_canvas = self.plot(dataset, new_train_state)
+        return new_canvas, new_train_state
+
+    def update_architecture(
+        self,
+        dataset: Dataset,
+        architecture: Architecture,
+        hyperparameters: Hyperparameters,
+    ):
+        print("Updating architecture")
+        new_model = init_model(architecture)
+        new_optimizer = init_optimizer(new_model, hyperparameters)
+        new_train_state = TrainState(new_model, new_optimizer)
+        new_canvas = self.plot(dataset, new_train_state)
+        return new_canvas, new_train_state
+
+    def update_hyperparameters(
+        self,
+        dataset: Dataset,
+        architecture: Architecture,
+        hyperparameters: Hyperparameters,
+    ):
+        print("Updating hyperparameters")
+        new_model = init_model(architecture)
+        new_optimizer = init_optimizer(new_model, hyperparameters)
+        new_train_state = TrainState(new_model, new_optimizer)
+        new_canvas = self.plot(dataset, new_train_state)
+        return new_canvas, new_train_state
+
+    def train_step(
+        self,
+        dataset: Dataset,
+        hyperparameters: Hyperparameters,
+        train_state: TrainState,
+    ):
+        print("Training step")
+        model = train_state.model
+        optimizer = train_state.optimizer
+        batch_size = hyperparameters.batch_size
+
+        model.train()
+        x_train = torch.from_numpy(dataset.x).float()
+        y_train = torch.from_numpy(dataset.y).float()
+
+        if batch_size < x_train.shape[0]:
+            indices = torch.randperm(x_train.shape[0])[:batch_size]
+            x_train = x_train[indices]
+            y_train = y_train[indices]
+
+        y_pred = model(x_train)
+        loss = nn.MSELoss()(y_pred.flatten(), y_train)
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        print(f"Training loss: {loss.item():.4f}")
+
+        new_canvas = self.plot(dataset, train_state)
+
+        return new_canvas, train_state
+
     def launch(self):
         # build the Gradio interface
         with gr.Blocks(css=self.css) as demo:
@@ -101,12 +234,17 @@ class MlpVisualizer:
             # states
             dataset = gr.State(Dataset())
             architecture = gr.State(Architecture())
+            hyperparameters = gr.State(Hyperparameters())
+
+            model = init_model(architecture.value)
+            optimizer = init_optimizer(model, hyperparameters.value)
+            train_state = gr.State(TrainState(model, optimizer))
 
             # GUI elements and layout 
             with gr.Row():
                 with gr.Column(scale=2):
                     canvas = gr.Image(
-                        value=self.plot(dataset.value, architecture.value),
+                        value=self.plot(dataset.value, train_state.value),
                         show_download_button=False,
                         container=True,
                     )
@@ -116,22 +254,37 @@ class MlpVisualizer:
                         dataset_view = DatasetView()
                         dataset_view.build(state=dataset)
                         dataset.change(
-                            fn=self.plot,
-                            inputs=[dataset],
-                            outputs=[canvas],
+                            fn=self.update_dataset,
+                            inputs=[dataset, architecture, hyperparameters],
+                            outputs=[canvas, train_state],
                         )
 
                     with gr.Tab("Architecture"):
                         architecture_view = ArchitectureView()
                         architecture_view.build(state=architecture)
                         architecture.change(
-                            fn=self.plot,
-                            inputs=[dataset, architecture],
-                            outputs=[canvas],
+                            fn=self.update_architecture,
+                            inputs=[dataset, architecture, hyperparameters],
+                            outputs=[canvas, train_state],
                         )
 
                     with gr.Tab("Train"):
-                        gr.Markdown("HI")
+                        hyperparameters_view = HyperparametersView()
+                        hyperparameters_view.build(state=hyperparameters)
+                        hyperparameters.change(
+                            fn=self.update_hyperparameters,
+                            inputs=[dataset, architecture, hyperparameters],
+                            outputs=[canvas, train_state],
+                        )
+
+                        train_button = gr.Button("Train 1 step")
+                        train_button.click(
+                            fn=self.train_step,
+                            inputs=[dataset, hyperparameters, train_state],
+                            outputs=[canvas, train_state],
+                        )
+
+
                     with gr.Tab("Plot"):
                         gr.Markdown("HI")
                     with gr.Tab("Export"):