mirror of https://github.com/hpcaitech/ColossalAI
39 lines
1.0 KiB
Python
39 lines
1.0 KiB
Python
#!/usr/bin/env python
|
|
# -*- encoding: utf-8 -*-
|
|
|
|
import torch.nn as nn
|
|
try:
|
|
import apex.amp as apex_amp
|
|
except ImportError:
|
|
pass
|
|
|
|
from torch import Tensor
|
|
|
|
from colossalai.nn.optimizer import ColossalaiOptimizer
|
|
from colossalai.utils import clip_grad_norm_fp32
|
|
|
|
|
|
class ApexAMPOptimizer(ColossalaiOptimizer):
|
|
""" A wrapper class for APEX optimizer and it implements apex-specific backward and clip_grad_norm
|
|
methods
|
|
"""
|
|
|
|
def backward(self, loss: Tensor):
|
|
"""Backward pass to get all gradients
|
|
|
|
Args:
|
|
loss (torch.Tensor): Loss computed by a loss function
|
|
"""
|
|
with apex_amp.scale_loss(loss, self.optim) as scaled_loss:
|
|
scaled_loss.backward()
|
|
|
|
def clip_grad_norm(self, model: nn.Module, max_norm: float):
|
|
"""Clip gradients by norm
|
|
|
|
Args:
|
|
model (torch.nn.Module): Your model object
|
|
max_norm (float): The max norm value for gradient clipping
|
|
"""
|
|
if max_norm > 0:
|
|
clip_grad_norm_fp32(apex_amp.master_params(self.optim), max_norm)
|