debug save_pretrained

Author: Tongjilibo

README.md

@@ -67,7 +67,7 @@ pip install git+https://github.com/Tongjilibo/torch4keras.git
 ## 4. 版本历史
 |更新日期| 版本 | 版本说明 |
 |------| ----------------- |----------- |
-|20240116|v0.1.8           | 重新整理snippets|
+|20240116|v0.1.8           | 重新整理snippets, 重写save_pretrained|
 |20231219|v0.1.7           | 增加SimpleStreamFileLogger和LoggerHandler, 修改Logger的格式|
 |20231208|v0.1.6.post2     |监控fit过程，有报错则发送邮件提醒; 解决torch2.0的compile冲突问题; 修复clip_grad_norm的bug|
 |20230928|v0.1.5     |进度条中显示已经训练的时间|

docs/History.md

@@ -1,6 +1,6 @@
 ## 更新历史
 
-- **20240116**: 重新整理snippets
+- **20240116**: 重新整理snippets, 重写save_pretrained
 - **20231219**: 增加SimpleStreamFileLogger和LoggerHandler, 修改Logger的格式
 - **20231208**: 监控fit过程，有报错则发送邮件提醒; 解决torch2.0的compile冲突问题
 - **20230928**: 进度条中显示已经训练的时间

examples/tutorials_mnist.py

@@ -76,7 +76,7 @@ def evaluate(self):
     email = EmailCallback(mail_receivers='[email protected]')  # 发送邮件
     wandb = WandbCallback(save_code=True)  # wandb
     hist = model.fit(train_dataloader, steps_per_epoch=steps_per_epoch, epochs=epochs, 
-                     callbacks=[Summary(), evaluator, logger, ckpt, early_stop])
+                     callbacks=[Summary(), evaluator, ts_board, logger, ckpt, early_stop])
 else:
     model.load_weights('./ckpt/5/model.pt')
     metrics = MyEvaluator().evaluate()

torch4keras/snippets/data_process.py

@@ -5,6 +5,8 @@
 from torch.utils.data import Dataset, IterableDataset
 import inspect
 from .import_utils import is_safetensors_available, is_sklearn_available
+import os
+
 
 if is_safetensors_available():
     from safetensors import safe_open
@@ -196,7 +198,7 @@ def metric_mapping(metric, func, y_pred, y_true):
     return None
 
 
-def load(checkpoint:str, load_safetensors:bool=False):
+def load_checkpoint(checkpoint:str, load_safetensors:bool=False):
     '''加载ckpt，支持torch.load和safetensors
     '''
     if load_safetensors or checkpoint.endswith(".safetensors"):
@@ -218,9 +220,12 @@ def load(checkpoint:str, load_safetensors:bool=False):
         return torch.load(checkpoint, map_location='cpu')
 
 
-def save(state_dict:dict, save_path:str, save_safetensors:bool=False):
+def save_checkpoint(state_dict:dict, save_path:str, save_safetensors:bool=False):
     '''保存ckpt，支持torch.save和safetensors
     '''
+    save_dir = os.path.dirname(save_path)
+    os.makedirs(save_dir, exist_ok=True)
+
     if save_safetensors or save_path.endswith('.safetensors'):
         safe_save_file(state_dict, save_path, metadata={"format": "pt"})
     else:

torch4keras/trainer.py

@@ -1,7 +1,7 @@
 from torch import nn
 import torch
 from torch4keras.snippets import DottableDict, metric_mapping, get_parameter_device, log_info, log_warn, log_error
-from torch4keras.snippets import print_trainable_parameters, colorful, monitor_run_by_email, load, save
+from torch4keras.snippets import print_trainable_parameters, colorful, monitor_run_by_email, load_checkpoint, save_checkpoint
 from torch4keras.callbacks import KerasProgbar, SmoothMetricsCallback, TqdmProgbar, ProgressBar2Progbar, Callback, CallbackList, History
 from collections import OrderedDict
 from typing import Union
@@ -460,7 +460,7 @@ def load_weights(self, load_path:Union[str,tuple,list], strict:bool=True, mappin
 
         mapping = mapping or dict()
         for load_path_i in load_path:
-            state_dict = load(load_path_i)
+            state_dict = load_checkpoint(load_path_i)
             for k in list(state_dict.keys()):
                 if k in mapping:
                     state_dict[mapping[k]] = state_dict.pop(k)
@@ -486,9 +486,7 @@ def save_weights(self, save_path:str, mapping:dict=None, trainable_only:bool=Fal
             if k in mapping:
                 state_dict[mapping[k]] = state_dict.pop(k)
 
-        save_dir = os.path.dirname(save_path)
-        os.makedirs(save_dir, exist_ok=True)
-        save(state_dict, save_path)
+        save_checkpoint(state_dict, save_path)
         if trainable_only:
             params_all = sum(p.numel() for p in self.unwrap_model().parameters())
             params_trainable = sum(p.numel() for p in self.unwrap_model().parameters() if p.requires_grad)
@@ -497,12 +495,15 @@ def save_weights(self, save_path:str, mapping:dict=None, trainable_only:bool=Fal
 
     def save_pretrained(self, save_path:str, weight_map:dict=None, mapping:dict=None):
         '''按照预训练模型的key来保存模型, 可供transformers包加载'''
+        state_dict = dict()
         for name, child in self.unwrap_model().named_children():
             if (name != '') and hasattr(child, 'save_pretrained'):
-                child.save_pretrained(save_path, weight_map, mapping)
+                tmp = child.save_pretrained(save_path, weight_map, mapping, write_to_disk=False)
+                state_dict.update(tmp if tmp else {})
             else:
-                save(child.state_dict(), save_path)
-
+                state_dict.update({f'{name}.{k}': v for k,v in child.state_dict().items()})
+        save_checkpoint(state_dict, save_path)
+    
     def resume_from_checkpoint(self, save_dir:str=None, model_path:str=None, optimizer_path:str=None, scheduler_path:str=None, 
                                steps_params_path:str=None, mapping:dict=None, verbose:int=0, strict:bool=True, **kwargs):
         '''同时加载模型、优化器、训练过程参数