Update models/videochat.py

models/videochat.py

@@ -143,19 +143,41 @@ class VideoChat(Blip2Base):
         for name, param in self.llama_model.named_parameters():
             param.requires_grad = False
         print('Loading LLAMA Done')
-
+        print(u'当前进程的内存使用：%.4f GB' % (psutil.Process(os.getpid()).memory_info().rss / 1024 / 1024 / 1024) )
+        info = psutil.virtual_memory()
+        print( u'电脑总内存：%.4f GB' % (info.total / 1024 / 1024 / 1024) )
+        print(u'当前使用的总内存占比：',info.percent)
+        print(u'cpu个数：',psutil.cpu_count())
         self.llama_proj = nn.Linear(
             self.Qformer.config.hidden_size, self.llama_model.config.hidden_size
         )
         self.max_txt_len = max_txt_len
-
+        print(u'当前进程的内存使用：%.4f GB' % (psutil.Process(os.getpid()).memory_info().rss / 1024 / 1024 / 1024) )
+        info = psutil.virtual_memory()
+        print( u'电脑总内存：%.4f GB' % (info.total / 1024 / 1024 / 1024) )
+        print(u'当前使用的总内存占比：',info.percent)
+        print(u'cpu个数：',psutil.cpu_count())
         # load weights of VideoChat
         if videochat_model_path:
+            print(u'当前进程的内存使用：%.4f GB' % (psutil.Process(os.getpid()).memory_info().rss / 1024 / 1024 / 1024) )
+            info = psutil.virtual_memory()
+            print( u'电脑总内存：%.4f GB' % (info.total / 1024 / 1024 / 1024) )
+            print(u'当前使用的总内存占比：',info.percent)
+            print(u'cpu个数：',psutil.cpu_count())
             print(f"Load VideoChat from: {videochat_model_path}")
             ckpt = torch.load(videochat_model_path, map_location="cpu")
+            print(u'ckpt load success.当前进程的内存使用：%.4f GB' % (psutil.Process(os.getpid()).memory_info().rss / 1024 / 1024 / 1024) )
+            info = psutil.virtual_memory()
+            print( u'电脑总内存：%.4f GB' % (info.total / 1024 / 1024 / 1024) )
+            print(u'当前使用的总内存占比：',info.percent)
+            print(u'cpu个数：',psutil.cpu_count())
             msg = self.load_state_dict(ckpt['model'], strict=False)
             print(msg)
-
+        print(u'当前进程的内存使用：%.4f GB' % (psutil.Process(os.getpid()).memory_info().rss / 1024 / 1024 / 1024) )
+        info = psutil.virtual_memory()
+        print( u'电脑总内存：%.4f GB' % (info.total / 1024 / 1024 / 1024) )
+        print(u'当前使用的总内存占比：',info.percent)
+        print(u'cpu个数：',psutil.cpu_count())
     def vit_to_cpu(self):
         self.ln_vision.to("cpu")
         self.ln_vision.float()