屏蔽torch.load警告, Timeit提示

Author: Tongjilibo

torch4keras/snippets/data_process.py

@@ -229,7 +229,10 @@ def load_checkpoint(checkpoint:str, load_safetensors:bool=False):
         return safe_load_file(checkpoint)
     else:
         # 正常加载pytorch_model.bin
-        return torch.load(checkpoint, map_location='cpu')
+        if 'weights_only' in inspect.signature(torch.load).parameters:
+            return torch.load(checkpoint, map_location='cpu', weights_only=True)
+        else:
+            return torch.load(checkpoint, map_location='cpu')
 
 
 def save_checkpoint(state_dict:dict, save_path:str, save_safetensors:bool=False):

torch4keras/snippets/monitor.py

@@ -152,6 +152,12 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         self.lap(name='Total', reset=False)
 
 
+class Timeit(TimeitContextManager):
+    def __init__(self) -> None:
+        super().__init__()
+        raise DeprecationWarning('`Timeit` has been deprecated since torch4keras==v0.2.8, use `TimeitContextManager` instead')
+
+
 class TimeitLogger:
     '''记录耗时
 
@@ -209,6 +215,12 @@ def end(self, verbose=1):
             return cost
 
 
+class Timeit2(TimeitLogger):
+    def __init__(self) -> None:
+        super().__init__()
+        raise DeprecationWarning('`Timeit2` has been deprecated since torch4keras==v0.2.8, use `TimeitLogger` instead')
+
+
 def send_email(mail_receivers:Union[str,list], mail_subject:str, mail_msg:str="", mail_host:str=None, 
                mail_user:str=None, mail_pwd:str=None, mail_sender:str=None):
     ''' 发送邮件(默认使用笔者自己注册的邮箱, 若含敏感信息请使用自己注册的邮箱)