Add log for memory usage tracking and fix bug with torch.set_grad_enabled

app.py

@@ -18,8 +18,6 @@ from factories import PhysicsWithGenerator, EvalModel, BaselineModel, EvalDatase
 ### Config
 # run model inference on NVIDIA gpu if available
 DEVICE_STR = 'cuda' if torch.cuda.is_available() else 'cpu'
-# stops tracking values for gradients
-torch.set_grad_enabled(False)
 
 
 ### Gradio Utils
@@ -62,17 +60,24 @@ def generate_imgs(x: torch.Tensor,
                   physics: PhysicsWithGenerator, use_gen: bool,
                   baseline: BaselineModel, model: EvalModel,
                   metrics: List[Metric]):
-    print(torch.cuda.memory_allocated() / 1024**2)
+    print(f"[Before inference] CUDA current allocated: {torch.cuda.memory_allocated() / 1024**2:.2f} MB")
+    print(f"[Before inference] CUDA current reserved: {torch.cuda.memory_reserved() / 1024**2:.2f} MB")
+    print(f"[Before inference] CUDA max allocated: {torch.cuda.max_memory_allocated() / 1024**2:.2f} MB")
+    print(f"[Before inference] CUDA max reserved: {torch.cuda.max_memory_reserved() / 1024**2:.2f} MB")
+
     ### Compute y
-    y = physics(x, use_gen)  # possible reduction in img shape due to Blurring
+    with torch.no_grad():
+        y = physics(x, use_gen)  # possible reduction in img shape due to Blurring
 
     ### Compute x_hat from RAM & DPIR
     ram_time = time.time()
-    out = model(y=y, physics=physics.physics)
+    with torch.no_grad():
+        out = model(y=y, physics=physics.physics)
     ram_time = time.time() - ram_time
 
     dpir_time = time.time()
-    out_baseline = baseline(y=y, physics=physics.physics)
+    with torch.no_grad():
+        out_baseline = baseline(y=y, physics=physics.physics)
     dpir_time = time.time() - dpir_time
 
     ### Process tensors before metric computation
@@ -87,13 +92,15 @@ def generate_imgs(x: torch.Tensor,
 
     ### Metrics
     metrics_y = ""
-    metrics_out = f"Inference time = {ram_time:.3f}s" + "\n"
-    metrics_out_baseline = f"Inference time = {dpir_time:.3f}s" + "\n"
+    metrics_out = ""
+    metrics_out_baseline = ""
     for metric in metrics:
         if y.shape == x.shape:
             metrics_y += f"{metric.name} = {metric(y, x).item():.4f}" + "\n"
         metrics_out += f"{metric.name} = {metric(out, x).item():.4f}" + "\n"
         metrics_out_baseline += f"{metric.name} = {metric(out_baseline, x).item():.4f}" + "\n"
+    metrics_out += f"Inference time = {ram_time:.3f}s"
+    metrics_out_baseline += f"Inference time = {dpir_time:.3f}s"
 
     ### Process y when y shape is different from x shape
     if physics.name == "MRI":
@@ -118,7 +125,10 @@ def generate_imgs(x: torch.Tensor,
     # Free memory
     del x, y, out, out_baseline, y_plot
     torch.cuda.empty_cache()
-    print(torch.cuda.memory_allocated() / 1024**2)
+    print(f"[After inference] CUDA current allocated: {torch.cuda.memory_allocated() / 1024**2:.2f} MB")
+    print(f"[After inference] CUDA current reserved: {torch.cuda.memory_reserved() / 1024**2:.2f} MB")
+    print(f"[After inference] CUDA max allocated: {torch.cuda.max_memory_allocated() / 1024**2:.2f} MB")
+    print(f"[After inference] CUDA max reserved: {torch.cuda.max_memory_reserved() / 1024**2:.2f} MB")
 
     return x_pil, y_pil, out_pil, out_baseline_pil, physics.display_saved_params(), metrics_y, metrics_out, metrics_out_baseline
 
@@ -150,13 +160,17 @@ def get_dataset(dataset_name):
 
 # global variables shared by all users
 ram_model = EvalModel("unext_emb_physics_config_C", device_str=DEVICE_STR)
+ram_model.eval()
 psnr = Metric.get_list_metrics(["PSNR"], device_str=DEVICE_STR)
 
 generate_imgs_from_user_partial = partial(generate_imgs_from_user, model=ram_model, metrics=psnr)
 generate_imgs_from_dataset_partial = partial(generate_imgs_from_dataset, model=ram_model, metrics=psnr)
 generate_random_imgs_from_dataset_partial = partial(generate_random_imgs_from_dataset, model=ram_model, metrics=psnr)
 
-print(torch.cuda.memory_allocated() / 1024**2)
+
+print(f"[Init] CUDA max allocated: {torch.cuda.max_memory_allocated() / 1024**2:.2f} MB")
+print(f"[Init] CUDA max reserved: {torch.cuda.max_memory_reserved() / 1024**2:.2f} MB")
+
 ### Gradio Blocks interface
 
 title = "Inverse problem playground"  # displayed on gradio tab and in the gradio page
@@ -172,7 +186,8 @@ with gr.Blocks(title=title, theme=gr.themes.Glass()) as interface:
     physics_placeholder = gr.State(lambda: get_physics_on_DEVICE_STR("MotionBlur_easy"))
     model_b_placeholder = gr.State(lambda: get_baseline_model_on_DEVICE_STR("DPIR"))
 
-    print(torch.cuda.memory_allocated() / 1024**2)
+    print(f"[Render] CUDA max allocated: {torch.cuda.max_memory_allocated() / 1024**2:.2f} MB")
+    print(f"[Render] CUDA max reserved: {torch.cuda.max_memory_reserved() / 1024**2:.2f} MB")
     @gr.render(inputs=[dataset_placeholder, physics_placeholder, available_physics_placeholder])
     def dynamic_layout(dataset, physics, available_physics):
         ### LAYOUT
@@ -196,11 +211,11 @@ with gr.Blocks(title=title, theme=gr.themes.Glass()) as interface:
                     load_button = gr.Button("Run on index image from dataset", size='md')
                     load_random_button = gr.Button("Run on random image from dataset", size='md')
             with gr.Column(scale=1, min_width=160):
-                observed_metrics = gr.Textbox(label="Observed metric", lines=3, key='metrics')
+                observed_metrics = gr.Textbox(label="Observed metric", lines=2, key='metrics')
             with gr.Column(scale=1, min_width=160):
-                out_a_metric = gr.Textbox(label="RAM output metrics", lines=3, key='ram_metrics')
+                out_a_metric = gr.Textbox(label="RAM output metrics", lines=2, key='ram_metrics')
             with gr.Column(scale=1, min_width=160):
-                out_b_metric = gr.Textbox(label="DPIR output metrics", lines=3, key='dpir_metrics')
+                out_b_metric = gr.Textbox(label="DPIR output metrics", lines=2, key='dpir_metrics')
 
         # Manage physics
         with gr.Row():