Add MLP architecture configuration

architecture.py

@@ -0,0 +1,148 @@
+import gradio as gr
+
+
+class Architecture:
+    def __init__(
+        self,
+        hidden_units: tuple[int] = (64, 64),
+        activations: tuple[str] = ("ReLU", "ReLU"),
+    ):
+        self.hidden_units = hidden_units
+        self.activations = activations
+
+    def update(self, **kwargs):
+        return Architecture(
+            hidden_units=kwargs.get("hidden_units", self.hidden_units),
+            activations=kwargs.get("activations", self.activations),
+        )
+
+    def __hash__(self):
+        return hash((self.hidden_units, self.activations))
+
+    @property
+    def num_layers(self):
+        return len(self.hidden_units)
+
+
+class ArchitectureView:
+    def __init__(self, max_layers: int = 5):
+        self.max_layers = max_layers
+
+    def update_layer_components(
+        self, state: Architecture, *layer_components
+    ):
+        if len(layer_components) != self.max_layers * 2:
+            raise ValueError("Incorrect number of layer components")
+
+        num_layers = state.num_layers
+
+        hidden_units = []
+        activations = []
+        for i in range(0, num_layers * 2, 2):
+            hidden_units.append(layer_components[i])
+            activations.append(layer_components[i + 1])
+
+        state = state.update(
+            hidden_units=tuple(hidden_units),
+            activations=tuple(activations),
+        )
+
+        return state
+
+    def add_layer(self, state: Architecture):
+        if state.num_layers < self.max_layers:
+            state = state.update(
+                hidden_units=state.hidden_units + (64,),
+                activations=state.activations + ("ReLU",),
+            )
+
+        updates = []
+        for i in range(self.max_layers):
+            # twice for hidden units and activation
+            updates.append(gr.update(visible=(i < state.num_layers)))
+            updates.append(gr.update(visible=(i < state.num_layers)))
+
+        return state, *updates
+
+    def remove_layer(self, state: Architecture):
+        if state.num_layers > 0:
+            state = state.update(
+                hidden_units=state.hidden_units[:-1],
+                activations=state.activations[:-1],
+            )
+
+        updates = []
+        for i in range(self.max_layers):
+            # twice for hidden units and activation
+            updates.append(gr.update(visible=(i < state.num_layers)))
+            updates.append(gr.update(visible=(i < state.num_layers)))
+
+        return state, *updates
+
+    def build(self, state: gr.State):
+        architecture = state.value
+
+        layer_components = []
+        with gr.Column():
+            with gr.Row():
+                add_layer = gr.Button("Add Layer")
+                remove_layer = gr.Button("Remove Layer")
+
+            for layer in range(self.max_layers):
+                with gr.Row():
+                    hidden_units = gr.Number(
+                        label="Hidden units",
+                        value=64,
+                        visible=(layer < architecture.num_layers),
+                        precision=0,
+                    )
+                    activation = gr.Dropdown(
+                        label="Activation",
+                        choices=["ReLU", "Sigmoid", "Tanh", "LeakyReLU", "ELU", "GELU", "Identity"],
+                        value="ReLU",
+                        visible=(layer < architecture.num_layers),
+                    )
+
+                layer_components.append(hidden_units)
+                layer_components.append(activation)
+
+            with gr.Row():
+                output_units = gr.Number(
+                    label="Output units",
+                    value=1,
+                    interactive=False,
+                )
+                output_activation = gr.Textbox(
+                    label="Activation",
+                    value="Identity",
+                    interactive=False,
+                )
+
+        # callbacks
+        add_layer.click(
+            fn=self.add_layer,
+            inputs=[state],
+            outputs=[state] + layer_components,
+        )
+        remove_layer.click(
+            fn=self.remove_layer,
+            inputs=[state],
+            outputs=[state] + layer_components,
+        )
+
+        for i, component in enumerate(layer_components):
+            # hidden unit
+            if i % 2 == 0:
+                component.submit(
+                    fn=self.update_layer_components,
+                    inputs=[state] + layer_components,
+                    outputs=[state],
+                )
+            # activation
+            else:
+                component.change(
+                    fn=self.update_layer_components,
+                    inputs=[state] + layer_components,
+                    outputs=[state],
+                )
+

mlp_visualizer.py

@@ -29,6 +29,7 @@ logging.basicConfig(
 )
 logger = logging.getLogger("ELVIS")
 
+from architecture import Architecture, ArchitectureView
 from dataset import Dataset, DatasetView, get_function
 
 
@@ -44,7 +45,7 @@ class MlpVisualizer:
     display: none;
 }"""
 
-    def plot(self, dataset_options: Dataset):
+    def plot(self, dataset: Dataset, architecture: Architecture) -> Image.Image:
         print("Plotting")
         t1 = time.time()
         fig = plt.figure(figsize=(self.canvas_width / 100., self.canvas_height / 100.0), dpi=100)
@@ -53,8 +54,8 @@ class MlpVisualizer:
         ax = fig.add_axes([0., 0., 1., 1.]) # 
         ax.margins(x=0, y=0) # no padding in both directions
 
-        if dataset_options.mode == "generate":
-            x_test, y_test = get_function(dataset_options.function, xlim=(-2, 2), nsample=100)
+        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()
 
@@ -64,15 +65,15 @@ class MlpVisualizer:
         ax.set_xlabel("x")
         ax.set_ylabel("y")
 
-        if dataset_options.mode == "generate":
+        if dataset.mode == "generate":
             ax.set_ylim(y_test.min() - 1, y_test.max() + 1)
 
-        x_train = dataset_options.x
-        y_train = dataset_options.y
+        x_train = dataset.x
+        y_train = dataset.y
         if True:
             plt.scatter(x_train.flatten(), y_train, label='training data', color=self.plot_cmap(0))
 
-        if dataset_options.mode == "generate":
+        if dataset.mode == "generate":
             plt.plot(x_test.flatten(), y_test, label='true function', color=self.plot_cmap(1))
 
         if False:
@@ -98,13 +99,14 @@ class MlpVisualizer:
             gr.HTML("<div style='text-align:left; font-size:40px; font-weight: bold;'>MLP Training Visualizer</div>")
 
             # states
-            dataset_options = gr.State(Dataset())
+            dataset = gr.State(Dataset())
+            architecture = gr.State(Architecture())
 
             # GUI elements and layout 
             with gr.Row():
                 with gr.Column(scale=2):
                     canvas = gr.Image(
-                        value=self.plot(dataset_options.value),
+                        value=self.plot(dataset.value, architecture.value),
                         show_download_button=False,
                         container=True,
                     )
@@ -112,15 +114,22 @@ class MlpVisualizer:
                 with gr.Column(scale=1): 
                     with gr.Tab("Dataset"):
                         dataset_view = DatasetView()
-                        dataset_view.build(state=dataset_options)
-                        dataset_options.change(
+                        dataset_view.build(state=dataset)
+                        dataset.change(
                             fn=self.plot,
-                            inputs=[dataset_options],
+                            inputs=[dataset],
                             outputs=[canvas],
                         )
 
                     with gr.Tab("Architecture"):
-                        gr.Markdown("HI")
+                        architecture_view = ArchitectureView()
+                        architecture_view.build(state=architecture)
+                        architecture.change(
+                            fn=self.plot,
+                            inputs=[dataset, architecture],
+                            outputs=[canvas],
+                        )
+
                     with gr.Tab("Train"):
                         gr.Markdown("HI")
                     with gr.Tab("Plot"):