mirrors/pytorch-image-models
1
0
Fork 0
mirror of https://github.com/huggingface/pytorch-image-models.git synced 2025-06-03 15:01:08 +08:00
Code Issues Projects Releases Wiki Activity

Clean up no_grad for trunc normal weight inits

Browse Source
This commit is contained in:
Ross Wightman 2022-08-25 16:29:52 -07:00
parent 48e1df8b37
commit 769ab4b98a
1 changed files with 25 additions and 23 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

48
timm/models/layers/weight_init.py
View File

@ -5,7 +5,7 @@ import warnings
from torch.nn.init import _calculate_fan_in_and_fan_out from torch.nn.init import _calculate_fan_in_and_fan_out
def _no_grad_trunc_normal_(tensor, mean, std, a, b): def _trunc_normal_(tensor, mean, std, a, b):
# Cut & paste from PyTorch official master until it's in a few official releases - RW # Cut & paste from PyTorch official master until it's in a few official releases - RW
# Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
def norm_cdf(x): def norm_cdf(x):
@ -17,28 +17,27 @@ def _no_grad_trunc_normal_(tensor, mean, std, a, b):
"The distribution of values may be incorrect.", "The distribution of values may be incorrect.",
stacklevel=2) stacklevel=2)
with torch.no_grad(): # Values are generated by using a truncated uniform distribution and
# Values are generated by using a truncated uniform distribution and # then using the inverse CDF for the normal distribution.
# then using the inverse CDF for the normal distribution. # Get upper and lower cdf values
# Get upper and lower cdf values l = norm_cdf((a - mean) / std)
l = norm_cdf((a - mean) / std) u = norm_cdf((b - mean) / std)
u = norm_cdf((b - mean) / std)
# Uniformly fill tensor with values from [l, u], then translate to # Uniformly fill tensor with values from [l, u], then translate to
# [2l-1, 2u-1]. # [2l-1, 2u-1].
tensor.uniform_(2 * l - 1, 2 * u - 1) tensor.uniform_(2 * l - 1, 2 * u - 1)
# Use inverse cdf transform for normal distribution to get truncated # Use inverse cdf transform for normal distribution to get truncated
# standard normal # standard normal
tensor.erfinv_() tensor.erfinv_()
# Transform to proper mean, std # Transform to proper mean, std
tensor.mul_(std * math.sqrt(2.)) tensor.mul_(std * math.sqrt(2.))
tensor.add_(mean) tensor.add_(mean)
# Clamp to ensure it's in the proper range # Clamp to ensure it's in the proper range
tensor.clamp_(min=a, max=b) tensor.clamp_(min=a, max=b)
return tensor return tensor
def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.): def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
@ -64,7 +63,8 @@ def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
>>> w = torch.empty(3, 5) >>> w = torch.empty(3, 5)
>>> nn.init.trunc_normal_(w) >>> nn.init.trunc_normal_(w)
""" """
return _no_grad_trunc_normal_(tensor, mean, std, a, b) with torch.no_grad():
return _trunc_normal_(tensor, mean, std, a, b)
def trunc_normal_tf_(tensor, mean=0., std=1., a=-2., b=2.): def trunc_normal_tf_(tensor, mean=0., std=1., a=-2., b=2.):
@ -90,8 +90,8 @@ def trunc_normal_tf_(tensor, mean=0., std=1., a=-2., b=2.):
>>> w = torch.empty(3, 5) >>> w = torch.empty(3, 5)
>>> nn.init.trunc_normal_(w) >>> nn.init.trunc_normal_(w)
""" """
_no_grad_trunc_normal_(tensor, 0, 1.0, a, b)
with torch.no_grad(): with torch.no_grad():
_trunc_normal_(tensor, 0, 1.0, a, b)
tensor.mul_(std).add_(mean) tensor.mul_(std).add_(mean)
return tensor return tensor
@ -111,10 +111,12 @@ def variance_scaling_(tensor, scale=1.0, mode='fan_in', distribution='normal'):
# constant is stddev of standard normal truncated to (-2, 2) # constant is stddev of standard normal truncated to (-2, 2)
trunc_normal_tf_(tensor, std=math.sqrt(variance) / .87962566103423978) trunc_normal_tf_(tensor, std=math.sqrt(variance) / .87962566103423978)
elif distribution == "normal": elif distribution == "normal":
tensor.normal_(std=math.sqrt(variance)) with torch.no_grad():
tensor.normal_(std=math.sqrt(variance))
elif distribution == "uniform": elif distribution == "uniform":
bound = math.sqrt(3 * variance) bound = math.sqrt(3 * variance)
tensor.uniform_(-bound, bound) with torch.no_grad():
tensor.uniform_(-bound, bound)
else: else:
raise ValueError(f"invalid distribution {distribution}") raise ValueError(f"invalid distribution {distribution}")