Update models/videochat.py

models/videochat.py

@@ -1,6 +1,7 @@
 import os
 import random
-import logging
+import
+ logging
 
 import torch
 from torch.cuda.amp import autocast as autocast
@@ -11,6 +12,8 @@ from .modeling_llama import LlamaForCausalLM
 from transformers import LlamaTokenizer, LlamaConfig
 
 
+
+
 class VideoChat(Blip2Base):
     """
     VideoChat model.
@@ -113,6 +116,14 @@ class VideoChat(Blip2Base):
         self.llama_tokenizer = LlamaTokenizer.from_pretrained(llama_model_path, use_fast=False, use_auth_token=os.environ["HF_TOKEN"])
         self.llama_tokenizer.pad_token = self.llama_tokenizer.eos_token
 
+        import psutil
+        import os
+        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())
+
         if self.low_resource:
                 self.llama_model = LlamaForCausalLM.from_pretrained(
                     llama_model_path,