Display current training loss

mlp_visualizer.py

@@ -195,10 +195,10 @@ class MlpVisualizer:
         # do not initialise here, otherwise gradio will make it not work
         # self.param_components = {}
 
-        self.model, self.optimizer = self.init_model()
+        self.criterion = nn.MSELoss()
+        self.model, self.optimizer, self.train_loss = self.init_model()
         self.num_steps_trained = 0
 
-        self.criterion = nn.MSELoss()
 
         self.plot_options = {
             "show_training_data": True,
@@ -272,7 +272,15 @@ class MlpVisualizer:
 
         self.num_steps_trained = 0
 
-        return model, optimizer
+        # compute initial train loss
+        model.eval()
+        inputs = torch.from_numpy(self.x_train).float()
+        targets = torch.from_numpy(self.y_train).float().unsqueeze(1)
+        with torch.no_grad():
+            outputs = model(inputs)
+            train_loss = self.criterion(outputs, targets).item()
+
+        return model, optimizer, train_loss
 
     def plot(self):
         '''
@@ -321,11 +329,11 @@ class MlpVisualizer:
         self.data_options["seed"] += 1
         self.x_train, self.y_train = self.generate_data()
         self.reset_model()
-        return self.plot(), self.num_steps_trained
+        return self.plot(), self.num_steps_trained, self.train_loss
 
     def reset_model(self):
-        self.model, self.optimizer = self.init_model()
-        return self.plot(), self.num_steps_trained
+        self.model, self.optimizer, self.train_loss = self.init_model()
+        return self.plot(), self.num_steps_trained, self.train_loss
 
     def update_data_options(self, **kwargs):
         for key, value in kwargs.items():
@@ -347,9 +355,9 @@ class MlpVisualizer:
 
         if "nsample" in kwargs:
             slider_update = gr.update(maximum=self.x_train.shape[0], value=min(self.basic_train_hparams["batch_size"], self.x_train.shape[0]))
-            return self.plot(), slider_update, self.num_steps_trained
+            return self.plot(), slider_update, self.num_steps_trained, self.train_loss
 
-        return self.plot(), self.num_steps_trained
+        return self.plot(), self.num_steps_trained, self.train_loss
 
     def update_plot_options(self, **kwargs):
         for key, value in kwargs.items():
@@ -362,9 +370,9 @@ class MlpVisualizer:
         self.architecture_options["activations"] = activations
 
         # reset model
-        self.model, self.optimizer = self.init_model()
+        self.model, self.optimizer, self.train_loss = self.init_model()
 
-        return self.plot(), self.num_steps_trained
+        return self.plot(), self.num_steps_trained, self.train_loss
 
     def update_basic_train_hparams(self, **kwargs):
         for key, value in kwargs.items():
@@ -372,9 +380,9 @@ class MlpVisualizer:
                 self.basic_train_hparams[key] = value
 
         # reset model
-        self.model, self.optimizer = self.init_model()
+        self.model, self.optimizer, self.train_loss = self.init_model()
 
-        return self.plot(), self.num_steps_trained
+        return self.plot(), self.num_steps_trained, self.train_loss
 
     def update_optimizer(self, optimizer_name):
         self.basic_train_hparams["optimizer"] = optimizer_name
@@ -388,9 +396,9 @@ class MlpVisualizer:
                 updates.append(gr.update(visible=is_visible))
 
         # reset model
-        self.model, self.optimizer = self.init_model()
+        self.model, self.optimizer, self.train_loss = self.init_model()
 
-        return updates + [self.plot(), self.num_steps_trained]
+        return updates + [self.plot(), self.num_steps_trained, self.train_loss]
 
     def build_optimizer_components(self):
         self.param_components = {}
@@ -414,8 +422,8 @@ class MlpVisualizer:
         self.optimizer_hparams[optimizer_name][param_name] = value
 
         # reset model and plot
-        self.model, self.optimizer = self.init_model()
-        return self.plot(), self.num_steps_trained
+        self.model, self.optimizer, self.train_loss = self.init_model()
+        return self.plot(), self.num_steps_trained, self.train_loss
     
     def train_step(self):
         self.model.train()
@@ -429,10 +437,15 @@ class MlpVisualizer:
         loss.backward()
         self.optimizer.step()
 
-        print(f"Training loss: {loss.item():.4f}")
         self.num_steps_trained += 1
 
-        return self.plot(), self.num_steps_trained
+        # update train loss
+        self.model.eval()
+        with torch.no_grad():
+            outputs = self.model(inputs)
+            self.train_loss = self.criterion(outputs, targets).item()
+
+        return self.plot(), self.num_steps_trained, self.train_loss
 
     def launch(self):
         # build the Gradio interface
@@ -525,6 +538,11 @@ class MlpVisualizer:
                                 value=0,
                                 interactive=False,
                             )
+                            train_loss_display = gr.Number(
+                                label="Train loss",
+                                value=self.train_loss,
+                                interactive=False,
+                            )
 
                         train_button = gr.Button("Train Step")
                         reset_model_button = gr.Button("Reset Model")
@@ -560,59 +578,59 @@ class MlpVisualizer:
                     function_box.submit(
                         fn=lambda function: self.update_data_options(function=function),
                         inputs=function_box,
-                        outputs=[self.canvas, train_step_counter],
+                        outputs=[self.canvas, train_step_counter, train_loss_display],
                     )
                     x_min.submit(
                         fn=lambda xmin: self.update_data_options(x_min=xmin),
                         inputs=x_min,
-                        outputs=[self.canvas, train_step_counter],
+                        outputs=[self.canvas, train_step_counter, train_loss_display],
                     )
                     x_max.submit(
                         fn=lambda xmax: self.update_data_options(x_max=xmax),
                         inputs=x_max,
-                        outputs=[self.canvas, train_step_counter],
+                        outputs=[self.canvas, train_step_counter, train_loss_display],
                     )
                     num_points_slider.change(
                         fn=lambda nsample: self.update_data_options(nsample=nsample),
                         inputs=num_points_slider,
-                        outputs=[self.canvas, batch_size_slider, train_step_counter],
+                        outputs=[self.canvas, batch_size_slider, train_step_counter, train_loss_display],
                     )
                     noise_value.submit(
                         fn=lambda sigma: self.update_data_options(sigma=sigma),
                         inputs=noise_value,
-                        outputs=[self.canvas, train_step_counter],
+                        outputs=[self.canvas, train_step_counter, train_loss_display],
                     )
                     regenerate_button.click(
                         fn=self._update_data_seed,
-                        outputs=[self.canvas, train_step_counter],
+                        outputs=[self.canvas, train_step_counter, train_loss_display],
                     )
             
                     # train options
                     optimizer_radio.change(
                         fn=self.update_optimizer,
                         inputs=optimizer_radio,
-                        outputs=[*all_param_components, self.canvas, train_step_counter],
+                        outputs=[*all_param_components, self.canvas, train_step_counter, train_loss_display],
                     )
                     batch_size_slider.change(
                         fn=lambda batch_size: self.update_basic_train_hparams(batch_size=batch_size),
                         inputs=batch_size_slider,
-                        outputs=[self.canvas, train_step_counter],
+                        outputs=[self.canvas, train_step_counter, train_loss_display],
                     )
                     train_button.click(
                         fn=self.train_step,
-                        outputs=[self.canvas, train_step_counter],
+                        outputs=[self.canvas, train_step_counter, train_loss_display],
                         show_progress="hidden",
                     )
                     reset_model_button.click(
                         fn=self.reset_model,
-                        outputs=[self.canvas, train_step_counter],
+                        outputs=[self.canvas, train_step_counter, train_loss_display],
                     )
                     for opt_name, params in self.param_components.items():
                         for param_name, comp in params.items():
                             comp.submit(
                                 fn=functools.partial(self.update_hparam, optimizer_name=opt_name, param_name=param_name),
                                 inputs=[comp],
-                                outputs=[self.canvas, train_step_counter],
+                                outputs=[self.canvas, train_step_counter, train_loss_display],
                             )
 
                     # plot options