mirror of
https://github.com/wassname/pytorch-lightning.git
synced 2026-09-21 13:20:08 +08:00
4c5e82c065
* Added default parser for trainer and class method to construct trainer from default args * Removed print statement * Added test for constructing Trainer from command line args * Removed extra line * Removed redundant imports, removed whitespace from empty lines * Fixed typo * Updated default parser creation to get class attributes automatically * Updated default parser creation to get class attributes automatically * Added method to get default args for trainer * Trimmed trainer get default args method * Updated from argparse method to not return trainer with static arguments * Update trainer get default args to classmethod * adjustment * fix * Fixed variable name * Update trainer.py * Update test_trainer.py * Update trainer.py * Update tests/trainer/test_trainer.py Co-Authored-By: Jirka Borovec <Borda@users.noreply.github.com> * Update trainer.py * Update test_trainer.py * Update trainer.py * Update test_trainer.py * Update tests/trainer/test_trainer.py Co-Authored-By: Jirka Borovec <Borda@users.noreply.github.com> * Update pytorch_lightning/trainer/trainer.py Co-Authored-By: Jirka Borovec <Borda@users.noreply.github.com> * Update trainer.py * Update test_trainer.py Co-authored-by: Mudit Tanwani <mudittanwani@gmail.com> Co-authored-by: Jirka Borovec <Borda@users.noreply.github.com>
62 lines
2.6 KiB
Python
62 lines
2.6 KiB
Python
import logging as log
|
|
from abc import ABC, abstractmethod
|
|
|
|
import torch
|
|
import math
|
|
|
|
from pytorch_lightning.callbacks import GradientAccumulationScheduler
|
|
|
|
EPSILON = 1e-6
|
|
EPSILON_FP16 = 1e-5
|
|
|
|
|
|
class TrainerTrainingTricksMixin(ABC):
|
|
|
|
# this is just a summary on variables used in this abstract class,
|
|
# the proper values/initialisation should be done in child class
|
|
gradient_clip_val: ...
|
|
|
|
@abstractmethod
|
|
def get_model(self):
|
|
"""Warning: this is just empty shell for code implemented in other class."""
|
|
|
|
def clip_gradients(self):
|
|
# this code is a modification of torch.nn.utils.clip_grad_norm_
|
|
# with TPU support based on https://github.com/pytorch/xla/blob/master/TROUBLESHOOTING.md
|
|
if self.gradient_clip_val > 0:
|
|
model = self.get_model()
|
|
parameters = model.parameters()
|
|
max_norm = float(self.gradient_clip_val)
|
|
norm_type = float(2.0)
|
|
if isinstance(parameters, torch.Tensor):
|
|
parameters = [parameters]
|
|
parameters = list(filter(lambda p: p.grad is not None, parameters))
|
|
if norm_type == math.inf:
|
|
total_norm = max(p.grad.data.abs().max() for p in parameters)
|
|
else:
|
|
device = parameters[0].device
|
|
total_norm = torch.zeros([], device=device if parameters else None)
|
|
for p in parameters:
|
|
param_norm = p.grad.data.norm(norm_type) ** norm_type
|
|
total_norm.add_(param_norm)
|
|
total_norm = (total_norm ** (1. / norm_type))
|
|
eps = EPSILON_FP16 if self.precision == 16 else EPSILON
|
|
clip_coef = torch.tensor(max_norm, device=device) / (total_norm + eps)
|
|
for p in parameters:
|
|
p.grad.data.mul_(torch.where(clip_coef < 1, clip_coef, torch.tensor(1., device=device)))
|
|
|
|
def print_nan_gradients(self):
|
|
model = self.get_model()
|
|
for param in model.parameters():
|
|
if (param.grad is not None) and torch.isnan(param.grad.float()).any():
|
|
log.info(param, param.grad)
|
|
|
|
def configure_accumulated_gradients(self, accumulate_grad_batches):
|
|
if isinstance(accumulate_grad_batches, dict):
|
|
self.accumulation_scheduler = GradientAccumulationScheduler(accumulate_grad_batches)
|
|
elif isinstance(accumulate_grad_batches, int):
|
|
schedule = {1: accumulate_grad_batches}
|
|
self.accumulation_scheduler = GradientAccumulationScheduler(schedule)
|
|
else:
|
|
raise TypeError("Gradient accumulation supports only int and dict types")
|