fix clip_grad_norm

Author: Tongjilibo

README.md

@@ -67,14 +67,14 @@ pip install git+https://github.com/Tongjilibo/torch4keras.git
 ## 4. 版本历史
 |更新日期| 版本 | 版本说明 |
 |------| ----------------- |----------- |
-|20231207| 0.1.6     |监控fit过程，有报错则发送邮件提醒; 解决torch2.0的compile冲突问题|
-|20230928| 0.1.5     |进度条中显示已经训练的时间|
+|20231208|v0.1.6.post2     |监控fit过程，有报错则发送邮件提醒; 解决torch2.0的compile冲突问题; 修复clip_grad_norm的bug|
+|20230928|v0.1.5     |进度条中显示已经训练的时间|
 |20230912|v0.1.4.post2|History增加plot()方法, 增加add_module()方法，修复0.1.4的_argparse_forward的bug, 增加loss2metrics方法|
 
 [更多版本](https://github.com/Tongjilibo/torch4keras/blob/master/docs/Update.md)
 
 ## 5. 更新历史：
-- **20231207**: 监控fit过程，有报错则发送邮件提醒; 解决torch2.0的compile冲突问题
+- **20231208**: 监控fit过程，有报错则发送邮件提醒; 解决torch2.0的compile冲突问题; 修复clip_grad_norm的bug
 - **20230928**: 进度条中显示已经训练的时间
 - **20230912**: History增加plot()方法, 增加add_module()方法，修复0.1.4的_argparse_forward的bug, 增加loss2metrics方法
 

docs/Update.md

@@ -2,8 +2,8 @@
 
 |更新日期| 版本 | 版本说明 |
 |------| ----------------- |----------- |
-|20231207| 0.1.6     |监控fit过程，有报错则发送邮件提醒; 解决torch2.0的compile冲突问题|
-|20230928| 0.1.5     |进度条中显示已经训练的时间|
+|20231207|v0.1.6     |监控fit过程，有报错则发送邮件提醒; 解决torch2.0的compile冲突问题|
+|20230928|v0.1.5     |进度条中显示已经训练的时间|
 |20230912|v0.1.4.post2|History增加plot()方法, 增加add_module()方法，修复0.1.4的_argparse_forward的bug, 增加loss2metrics方法|
 |20230909|v0.1.4|增加from_pretrained和save_pretrained方法，增加log_warn_once方法，compile()中可设置成员变量，默认move_to_model_device设置为True, 增加JsonConfig，增加_argparse_forward()方便下游继承改写Trainer|
 |20230901|v0.1.3|compile()可不传参，interval不一致报warning, 去除部分self.vars, 调整move_to_model_device逻辑，DDP每个epoch重新设置随机数，save_weights()和load_weights()可以按照`pretrained`格式|

examples/tutorials_mnist.py

@@ -54,7 +54,7 @@
 model = Trainer(net.to(device))
 optimizer = optim.Adam(net.parameters())
 scheduler = get_linear_schedule_with_warmup(optimizer, steps_per_epoch, steps_per_epoch*epochs)
-model.compile(optimizer=optimizer, scheduler=scheduler, loss=nn.CrossEntropyLoss(), metrics=['acc'])
+model.compile(optimizer=optimizer, scheduler=scheduler, loss=nn.CrossEntropyLoss(), metrics=['acc'], clip_grad_norm=1.0)
 
 class MyEvaluator(Evaluator):
     # 重构评价函数

setup.py

@@ -7,7 +7,7 @@
 
 setup(
     name='torch4keras',
-    version='v0.1.6',
+    version='v0.1.6.post2',
     description='Use torch like keras',
     long_description=long_description,
     long_description_content_type="text/markdown",

torch4keras/trainer.py

@@ -217,13 +217,13 @@ def step(self):
         if self.mixed_precision:
             self.scaler.unscale_(self.optimizer)
             if self.clip_grad_norm is not None:  # 梯度裁剪
-                torch.nn.utils.clip_grad_norm_(self.parameters(), self.clip_grad_norm)
+                torch.nn.utils.clip_grad_norm_(self.unwrap_model().parameters(), self.clip_grad_norm)
             self.scaler.step(self.optimizer)
             self.scaler.update()
             skip_scheduler = self.scaler.get_scale() != self.scale_before_step
         else:
             if self.clip_grad_norm is not None:  # 梯度裁剪
-                torch.nn.utils.clip_grad_norm_(self.parameters(), self.clip_grad_norm)
+                torch.nn.utils.clip_grad_norm_(self.unwrap_model().parameters(), self.clip_grad_norm)
             self.optimizer.step()
 
         self.optimizer.zero_grad()  # 清梯度