init

app.py

@@ -35,21 +35,23 @@ def make_dataframe(x, y, method=None):
         df = pd.DataFrame({'x': x, 'y': y})
     return df
 
+
 Y_demo = torch.linspace(Y_LB, Y_UB, 2).unsqueeze(-1)
 X_demo = (Y_demo - B) / K
 
 df_oracle = make_dataframe(X_demo, Y_demo, 'Oracle')
 
+
 def prepare_data(sel_num):
     interval = (Y_UB - Y_LB) / NUM_SEG
     all_x, all_y = [], []
     prob = []
     for i in range(NUM_SEG):
-        uniform_y_distribution = torch.distributions.Uniform(Y_UB - (i+1)*interval, Y_UB-i*interval)
-        y_uniform = uniform_y_distribution.sample((sel_num[i], 1))
+        uniform_y_distribution = torch.distributions.Uniform(Y_UB - (i + 1) * interval, Y_UB - i * interval)
+        y_uniform = uniform_y_distribution.sample((NUM_TRAIN_SAMPLES, 1))[:sel_num[i]]
 
         noise_distribution = torch.distributions.Normal(loc=0, scale=NOISE_SIGMA)
-        noise = noise_distribution.sample((sel_num[i], 1))
+        noise = noise_distribution.sample((NUM_TRAIN_SAMPLES, 1))[:sel_num[i]]
         y_uniform_oracle = y_uniform - noise
 
         x_uniform = (y_uniform_oracle - B) / K
@@ -73,6 +75,7 @@ def unzip_dataloader(training_loader):
     all_y = torch.cat(all_y)
     return all_x, all_y
 
+
 def train(train_loader, training_bundle, num_epochs):
     training_df = make_dataframe(*unzip_dataloader(train_loader))
     for epoch in range(num_epochs):
@@ -91,6 +94,7 @@ def train(train_loader, training_bundle, num_epochs):
         if (epoch + 1) % PRINT_FREQ == 0:
             visualize(training_df, training_bundle, epoch)
 
+
 def visualize(training_df, training_bundle, epoch):
     df = df_oracle
     for model, optimizer, scheduler, criterion, criterion_name in training_bundle:
@@ -103,11 +107,10 @@ def visualize(training_df, training_bundle, epoch):
     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.savefig('train_log/{:05d}.png'.format(epoch + 1), bbox_inches='tight')
     plt.close()
 
 
-
 def make_video():
     (
         ffmpeg
@@ -116,6 +119,7 @@ def make_video():
             .run()
     )
 
+
 class ReweightL2(_Loss):
     def __init__(self, reweight='inverse'):
         super(ReweightL2, self).__init__()
@@ -134,6 +138,7 @@ class ReweightL2(_Loss):
         loss = loss.sum()
         return loss
 
+
 class LinearModel(nn.Module):
     def __init__(self, input_dim, output_dim):
         super(LinearModel, self).__init__()
@@ -145,6 +150,7 @@ class LinearModel(nn.Module):
         x = self.mlp(x)
         return x
 
+
 def prepare_model():
     model = LinearModel(input_dim=1, output_dim=1)
     optimizer = torch.optim.SGD(model.parameters(), lr=1e-2, momentum=0.9)
@@ -171,6 +177,7 @@ def bmc_loss(pred, target, noise_var):
 
     return loss * (2 * noise_var)
 
+
 def regress(train_loader):
     training_bundle = []
     criterions = {
@@ -184,6 +191,7 @@ def regress(train_loader):
         training_bundle.append((model, optimizer, scheduler, criterion, criterion_name))
     train(train_loader, training_bundle, NUM_EPOCHS)
 
+
 class DummyDataset(Dataset):
     def __init__(self, inputs, targets, prob):
         self.inputs = inputs
@@ -200,20 +208,21 @@ class DummyDataset(Dataset):
 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),
+                      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)
     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.savefig('training_data.png', bbox_inches='tight')
     plt.close()
 
+
 def clean_up_logs():
     if not osp.exists('train_log'):
         os.mkdir('train_log')
@@ -222,9 +231,10 @@ def clean_up_logs():
     if osp.isfile('movie.mp4'):
         os.remove('movie.mp4')
 
+
 def run(num1, num2, num3, num4, num5, random_seed, submit):
     sel_num = [num1, num2, num3, num4, num5]
-    sel_num = [int(num/100*NUM_PER_BUCKET) for num in sel_num]
+    sel_num = [int(num / 100 * NUM_PER_BUCKET) for num in sel_num]
     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)
@@ -238,25 +248,24 @@ def run(num1, num2, num3, num4, num5, random_seed, submit):
         clean_up_logs()
         regress(train_loader)
         make_video()
-    output = 'train_log/{:05d}.png'.format(NUM_EPOCHS) if submit==1 else None
-    video = "movie.mp4" if submit==1 else None
+    output = 'train_log/{:05d}.png'.format(NUM_EPOCHS) if submit == 1 else None
+    video = "movie.mp4" if submit == 1 else None
     return 'training_data.png', text, output, video
 
 
 if __name__ == '__main__':
-    
     iface = gr.Interface(
         fn=run,
         inputs=[
-                gr.inputs.Slider(0, 100, default=50, step=0.1, label='Label percentage in [0, 2)'),
-                gr.inputs.Slider(0, 100, default=50, step=0.1, label='Label percentage in [2, 4)'),
-                gr.inputs.Slider(0, 100, default=50, step=0.1, label='Label percentage in [4, 6)'),
-                gr.inputs.Slider(0, 100, default=50, step=0.1, label='Label percentage in [6, 8)'),
-                gr.inputs.Slider(0, 100, default=50, step=0.1, label='Label percentage in [8, 10)'),
-                gr.inputs.Number(default=0, label='Random Seed', optional=False),
-                gr.inputs.Radio(['Prepare Training Data', 'Start Regressing!'],
-                                type="index", default=None, label='Mode', optional=False),
-                ],
+            gr.inputs.Slider(0, 100, default=50, step=0.1, label='Label percentage in [0, 2)'),
+            gr.inputs.Slider(0, 100, default=50, step=0.1, label='Label percentage in [2, 4)'),
+            gr.inputs.Slider(0, 100, default=50, step=0.1, label='Label percentage in [4, 6)'),
+            gr.inputs.Slider(0, 100, default=50, step=0.1, label='Label percentage in [6, 8)'),
+            gr.inputs.Slider(0, 100, default=50, step=0.1, label='Label percentage in [8, 10)'),
+            gr.inputs.Number(default=0, label='Random Seed', optional=False),
+            gr.inputs.Radio(['Prepare Training Data', 'Start Regressing!'],
+                            type="index", default=None, label='Mode', optional=False),
+        ],
         outputs=[
             gr.outputs.Image(type="file", label="Training data"),
             gr.outputs.Textbox(type="auto", label='What\' s next?'),
@@ -273,6 +282,6 @@ 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'],
         ],
-        
+
     )
     iface.launch()