init

app.py

@@ -18,7 +18,6 @@ Y_LB = 0
 K = 1
 B = 0
 NUM_SEG = 5
-sns.set_theme(palette='colorblind')
 NUM_EPOCHS = 100
 PRINT_FREQ = NUM_EPOCHS // 20
 NUM_TRAIN_SAMPLES = NUM_PER_BUCKET * NUM_SEG
@@ -76,8 +75,8 @@ def unzip_dataloader(training_loader):
     return all_x, all_y
 
 
-def train(train_loader, training_bundle, num_epochs):
-    training_df = make_dataframe(*unzip_dataloader(train_loader))
+def train(train_loader, training_df, training_bundle, num_epochs):
+    visualize_training_process(training_df, training_bundle, -1)
     for epoch in range(num_epochs):
         for model, optimizer, scheduler, criterion, criterion_name in training_bundle:
             model.train()
@@ -92,23 +91,19 @@ def train(train_loader, training_bundle, num_epochs):
                 optimizer.step()
             scheduler.step()
         if (epoch + 1) % PRINT_FREQ == 0:
-            visualize(training_df, training_bundle, epoch)
+            visualize_training_process(training_df, training_bundle, epoch)
+    visualize_training_process(training_df, training_bundle, num_epochs, final=True)
 
 
-def visualize(training_df, training_bundle, epoch):
+def visualize_training_process(training_df, training_bundle, epoch, final=False):
     df = df_oracle
     for model, optimizer, scheduler, criterion, criterion_name in training_bundle:
         model.eval()
         y = model(X_demo)
         df = df.append(make_dataframe(X_demo, y, criterion_name), ignore_index=True)
-    sns.lineplot(data=df, x='x', y='y', hue='Method', estimator=None, ci=None)
-    sns.scatterplot(data=training_df, x='x', y='y', color='#003ea1', alpha=0.05, linewidths=0, s=100)
-    plt.xlim((Y_LB - B) / K, (Y_UB - B) / K)
-    plt.ylim(Y_LB, Y_UB)
-    plt.gca().axes.set_xlabel(r'$x$', fontsize=10)
-    plt.gca().axes.set_ylabel(r'$y$', fontsize=10)
-    plt.savefig('train_log/{:05d}.png'.format(epoch + 1), bbox_inches='tight')
-    plt.close()
+    visualize(training_df, df, 'train_log/{:05d}.png'.format(epoch + 1), fast=True)
+    if final:
+        visualize(training_df, df, 'regression_result.png', fast=False)
 
 
 def make_video():
@@ -178,7 +173,7 @@ def bmc_loss(pred, target, noise_var):
     return loss * (2 * noise_var)
 
 
-def regress(train_loader):
+def regress(train_loader, training_df):
     training_bundle = []
     criterions = {
         'MSE': torch.nn.MSELoss(),
@@ -189,7 +184,7 @@ def regress(train_loader):
         criterion = criterions[criterion_name]
         model, optimizer, scheduler = prepare_model()
         training_bundle.append((model, optimizer, scheduler, criterion, criterion_name))
-    train(train_loader, training_bundle, NUM_EPOCHS)
+    train(train_loader, training_df, training_bundle, NUM_EPOCHS)
 
 
 class DummyDataset(Dataset):
@@ -205,22 +200,31 @@ class DummyDataset(Dataset):
         return len(self.inputs)
 
 
-def vis_training_data(all_x, all_y):
-    training_df = make_dataframe(all_x, all_y)
-    g = sns.jointplot(data=training_df, x='x', y='y', color='#003ea1', alpha=0.1, linewidths=0, s=100,
-                      marginal_kws=dict(bins=torch.linspace(Y_LB, Y_UB, steps=NUM_SEG + 1), rug=True),
-                      xlim=((Y_LB - B) / K, (Y_UB - B) / K),
-                      ylim=(Y_LB, Y_UB),
-                      space=0.1,
-                      height=8,
-                      ratio=2
-                      )
-    g.ax_marg_x.remove()
-    sns.lineplot(data=df_oracle, x='x', y='y', hue='Method', ax=g.ax_joint, legend=False)
+def visualize(training_df, df, save_path, fast=False):
+    if fast:
+        g_line = sns.lineplot(data=df, x='x', y='y', hue='Method', estimator=None, ci=None)
+        plt.xlim((Y_LB - B) / K, (Y_UB - B) / K)
+        plt.ylim(Y_LB, Y_UB)
+    else:
+        g = sns.jointplot(data=training_df, x='x', y='y', color='#003ea1', alpha=0.1, linewidths=0, s=100,
+                          marginal_kws=dict(bins=torch.linspace(Y_LB, Y_UB, steps=NUM_SEG + 1)),
+                          xlim=((Y_LB - B) / K, (Y_UB - B) / K),
+                          ylim=(Y_LB, Y_UB),
+                          space=0.1,
+                          height=8,
+                          ratio=2,
+                          estimator=None, ci=None,
+                          legend=False
+                          )
+        g.ax_marg_x.remove()
+        g_line = sns.lineplot(data=df, x='x', y='y', hue='Method', ax=g.ax_joint, estimator=None, ci=None)
+    g_line.legend_.set_title(None)
+    g_line.legend(loc='upper left')
     plt.gca().axes.set_xlabel(r'$x$', fontsize=10)
     plt.gca().axes.set_ylabel(r'$y$', fontsize=10)
-    plt.savefig('training_data.png', bbox_inches='tight')
-    plt.close()
+
+    plt.savefig(save_path, bbox_inches='tight')
+    plt.clf()
 
 
 def clean_up_logs():
@@ -228,8 +232,9 @@ def clean_up_logs():
         os.mkdir('train_log')
     for f in os.listdir('train_log'):
         os.remove(osp.join('train_log', f))
-    if osp.isfile('movie.mp4'):
-        os.remove('movie.mp4')
+    for f in ['regression_result.png', 'training_data.png', 'movie.mp4']:
+        if osp.isfile(f):
+            os.remove(f)
 
 
 def run(num1, num2, num3, num4, num5, random_seed, submit):
@@ -238,19 +243,22 @@ def run(num1, num2, num3, num4, num5, random_seed, submit):
     torch.manual_seed(int(random_seed))
     all_x, all_y, prob = prepare_data(sel_num)
     train_loader = DataLoader(DummyDataset(all_x, all_y, prob), BATCH_SIZE, shuffle=True)
-    vis_training_data(all_x, all_y)
+    training_df = make_dataframe(all_x, all_y)
 
+    clean_up_logs()
     if submit == 0:
-        text = "Press \"Start Regressing\" if your are happy with the training data. Regression takes ~10s."
-    else:
-        text = "Press \"Prepare Training Data\" before changing the training data. You may also change the random seed."
+        visualize(training_df, df_oracle, 'training_data.png')
     if submit == 1:
-        clean_up_logs()
-        regress(train_loader)
+        regress(train_loader, training_df)
         make_video()
-    output = 'train_log/{:05d}.png'.format(NUM_EPOCHS) if submit == 1 else None
+    if submit == 0:
+        text = "Press \"Start Regressing\" if your are happy with the training data. Regression takes ~10s."
+    else:
+        text = "Press \"Prepare Training Data\" before moving the sliders. You may also change the random seed."
+    training_data_plot = 'training_data.png' if submit == 0 else None
+    output = 'regression_result.png'.format(NUM_EPOCHS) if submit == 1 else None
     video = "movie.mp4" if submit == 1 else None
-    return 'training_data.png', text, output, video
+    return training_data_plot, output, video, text
 
 
 if __name__ == '__main__':
@@ -268,9 +276,9 @@ if __name__ == '__main__':
         ],
         outputs=[
             gr.outputs.Image(type="file", label="Training data"),
-            gr.outputs.Textbox(type="auto", label='What\' s next?'),
             gr.outputs.Image(type="file", label="Regression result"),
-            gr.outputs.Video(type='mp4', label='Training process')
+            gr.outputs.Video(type='mp4', label='Training process'),
+            gr.outputs.Textbox(type="auto", label='What\' s next?')
         ],
         live=True,
         allow_flagging='never',
@@ -282,6 +290,7 @@ if __name__ == '__main__':
             [0.1, 0.8, 6.4, 51.2, 100, 0, 'Prepare Training Data'],
             [1, 10, 100, 10, 1, 0, 'Prepare Training Data'],
         ],
-
+        css = ".output-image, .image-preview {height: 500px !important}",
+        article="<p style='text-align: center'><a href='https://github.com/jiawei-ren/BalancedMSE' target='_blank'>Balanced MSE @ GitHub</a></p> "
     )
     iface.launch()