Update app.py

app.py

@@ -12,6 +12,10 @@ import torch.nn as nn
 import torch.nn.functional as F
 
 
+def log_gpu_memory():
+    print(subprocess.check_output('nvidia-smi').decode('utf-8'))
+    
+
 class RelationModuleMultiScale(torch.nn.Module):
 
     def __init__(self, img_feature_dim, num_bottleneck, num_frames):
@@ -264,11 +268,15 @@ def MyPlot(frame_id, src_orig, tar_orig, src_recon, tar_recon, src_Zt, tar_Zt, s
     plt.savefig(save_name, dpi=200, format='png', bbox_inches='tight', pad_inches=0.0)
     
 
+log_gpu_memory()
+
 # == Load Model ==
 model = TransferVAE_Video()
 model.load_state_dict(torch.load('TransferVAE.pth.tar', map_location=torch.device('cpu'))['state_dict'])
 model.eval()
-    
+
+log_gpu_memory()
+
   
 def run(source, action_source, hair_source, top_source, bottom_source, target, action_target, hair_target, top_target, bottom_target):
 
@@ -330,10 +338,12 @@ def run(source, action_source, hair_source, top_source, bottom_source, target, a
     x = torch.cat((images_source, images_target), dim=0)
     
     
+    log_gpu_memory()
     # == Forward ==
     with torch.no_grad():
         f_post, z_post, recon_x = model(x, [0]*3)
-     
+    log_gpu_memory()
+    
     src_orig_sample = x[0, :, :, :, :]
     src_recon_sample = recon_x[0, :, :, :, :]
     src_f_post = f_post[0, :].unsqueeze(0)
@@ -378,7 +388,9 @@ def run(source, action_source, hair_source, top_source, bottom_source, target, a
         recon_x_tarZf_srcZt = model.decoder_frame(zf_tarZf_srcZt)
         tar_Zf_src_Zt = recon_x_tarZf_srcZt.squeeze()[frame, :, :, :].detach().numpy().transpose((1, 2, 0))
 
+        log_gpu_memory()
         MyPlot(frame, src_orig, tar_orig, src_recon, tar_recon, src_Zt, tar_Zt, src_Zf_tar_Zt, tar_Zf_src_Zt)
+        log_gpu_memory()
         
     a = concat('MyPlot_')