[master] add grad_clip

Author: vacancy

examples/torch-starter/trainval.py

@@ -112,17 +112,6 @@ def main():
         else:
             args.tb_dir = None
 
-        mldash.init(
-            desc_name=args.series_name + '/' + args.desc_name,
-            expr_name=args.expr,
-            run_name=args.run_name,
-            args=args,
-            highlight_args=parser,
-            configs=configs,
-        )
-
-        mldash.update(metainfo_file=args.meta_file, log_file=args.log_file, meter_file=args.meter_file, tb_dir=args.tb_dir)
-
     if not args.debug:
         logger.critical('Writing logs to file: "{}".'.format(args.log_file))
         set_output_file(args.log_file)
@@ -194,8 +183,22 @@ def main():
         meters = GroupMeters()
 
     if not args.debug:
+        logger.critical('Writing metainfo to file: "{}".'.format(args.meta_file))
+        with open(args.meta_file, 'w') as f:
+            f.write(dump_metainfo(args=args.__dict__, configs=configs))
         logger.critical('Writing meter logs to file: "{}".'.format(args.meter_file))
 
+        logger.critical('Initializing MLDash.')
+        mldash.init(
+            desc_name=args.series_name + '/' + args.desc_name,
+            expr_name=args.expr,
+            run_name=args.run_name,
+            args=args,
+            highlight_args=parser,
+            configs=configs,
+        )
+        mldash.update(metainfo_file=args.meta_file, log_file=args.log_file, meter_file=args.meter_file, tb_dir=args.tb_dir)
+
     if args.embed:
         from IPython import embed; embed()
 

jactorch/train/env.py

@@ -116,7 +116,7 @@ def decay_learning_rate(self, decay):
         for param_group in self._optimizer.param_groups:
             param_group['lr'] *= decay
 
-    def step(self, feed_dict, reduce_func=default_reduce_func, cast_tensor=False, measure_time=False):
+    def step(self, feed_dict, grad_clip=0., reduce_func=default_reduce_func, cast_tensor=False, measure_time=False):
         if hasattr(self.model, 'train_step'):
             return self.model.train_step(self.optimizer, feed_dict)
 
@@ -153,6 +153,9 @@ def step(self, feed_dict, reduce_func=default_reduce_func, cast_tensor=False, me
         self.trigger_event('backward:before', self, feed_dict, loss, monitors, output_dict)
         if loss.requires_grad:
             loss.backward()
+            if grad_clip > 0:
+                from torch.nn.utils.clip_grad import clip_grad_norm_
+                clip_grad_norm_(self.model.parameters(), grad_clip)
 
         if measure_time:
             extra['time/backward'] = cuda_time() - end_time