update reset

app.py

@@ -151,36 +151,55 @@ def draw_arrow(ax, start, end):
                 ls="-", lw=1)
 
 def draw_reset_result(num_data=16, num_code=12):
-    fig, ax = plt.subplots(1, 6, figsize=(22, 4))
+    fig_reset, ax_reset = plt.subplots(1, 6, figsize=(36, 6), dpi=400)
+    fig_nreset, ax_nreset = plt.subplots(1, 6, figsize=(36, 6), dpi=400)
     x = torch.randn(num_data, 1) * 2 + 5
     y = torch.randn(num_data, 1) * 2 - 5
     data = torch.cat([x, y], dim=1)
     quantizer = Quantizer(TYPE='vq', code_dim=2, num_code=num_code, num_group=1, tokens_per_data=1)
     optimizer = torch.optim.SGD(quantizer.parameters(), lr=0.1)
-    draw_fig(ax[0], quantizer, data, title=f"Initialization")
-    ax[0].legend(["Data", "Code"], loc="upper right", fontsize=18)
+    quantizer_nreset = Quantizer(TYPE='vq', code_dim=2, num_code=num_code, num_group=1, tokens_per_data=1, auto_reset=False)
+    optimizer_nreset = torch.optim.SGD(quantizer_nreset.parameters(), lr=0.1)
+    draw_fig(ax_reset[0], quantizer, data, title=f"Initialization")
+    draw_fig(ax_nreset[0], quantizer_nreset, data, title=f"Initialization")
+    ax_reset[0].legend(["Data", "Code"], loc="upper right", fontsize=18)
+    ax_nreset[0].legend(["Data", "Code"], loc="upper right", fontsize=18)
 
-    i_list = [1, 2, 3, 10, 100]
+    i_list = [1, 3, 10, 50, 200]
 
     count = 0
     for i in range(1000):
         optimizer.zero_grad()
+        optimizer_nreset.zero_grad()
         output_dict = quantizer(data.unsqueeze(1))
+        output_dict_nreset = quantizer_nreset(data.unsqueeze(1))
         quant_data = output_dict["x_quant"].squeeze()
+        quant_data_nreset = output_dict_nreset["x_quant"].squeeze()
         indices = output_dict["indices"].squeeze()
+        indices = output_dict_nreset["indices"].squeeze()
         loss = torch.mean((quant_data - data) ** 2)
+        loss_nreset = torch.mean((quant_data_nreset - data) ** 2)
         loss.backward()
+        loss_nreset.backward()
         optimizer.step()
+        optimizer_nreset.step()
 
         if (i+1) in i_list:
             count += 1
-            draw_fig(ax[count], quantizer, data, title=f"Iters: {i+1}, MSE: {loss.item():.1f}")
-            draw_arrow(ax[count], quant_data.detach().numpy(), data.numpy())
+            draw_fig(ax_reset[count], quantizer, data, title=f"Iters: {i+1}, MSE: {loss.item():.1f}")
+            draw_arrow(ax_reset[count], quant_data.detach().numpy(), data.numpy())
+
+            draw_fig(ax_nreset[count], quantizer_nreset, data, title=f"Iters: {i+1}, MSE: {loss_nreset.item():.1f}")
+            draw_arrow(ax_nreset[count], quant_data_nreset.detach().numpy(), data.numpy())
 
         quantizer.reset()
 
-    return fig
+    fig_reset.suptitle("VQ Codebook Training with Reset", fontsize=24, y=1.05)
+    fig_nreset.suptitle("VQ Codebook Training without Reset", fontsize=24, y=1.05)
+
+    return fig_reset, fig_nreset
 
+# end
 
 
 class Handler:
@@ -290,7 +309,10 @@ if __name__ == "__main__":
             gr.Slider(label="num_data", value=16, minimum=10, maximum=20, step=1),
             gr.Slider(label="num_code", value=12, minimum=8, maximum=16, step=1),
         ],
-        outputs=gr.Plot(label="Training Visualization"),
+        outputs=[
+            gr.Plot(label="With Reset"),
+            gr.Plot(label="Without Reset")
+        ],
         title="Demo 2: Codebook Reset Strategy Visualization",
         description="Visualizes codebook and data movement at different training steps."
     )