| 42 | |
| 43 | |
| 44 | class DistOptimizerHook(OptimizerHook): |
| 45 | |
| 46 | def __init__(self, grad_clip=None, coalesce=True, bucket_size_mb=-1): |
| 47 | self.grad_clip = grad_clip |
| 48 | self.coalesce = coalesce |
| 49 | self.bucket_size_mb = bucket_size_mb |
| 50 | |
| 51 | def after_train_iter(self, runner): |
| 52 | runner.optimizer.zero_grad() |
| 53 | runner.outputs['loss'].backward() |
| 54 | if self.grad_clip is not None: |
| 55 | self.clip_grads(runner.model.parameters()) |
| 56 | runner.optimizer.step() |
nothing calls this directly
no outgoing calls
no test coverage detected