mirror of https://github.com/Jittor/Jittor
128 lines
3.8 KiB
Python
128 lines
3.8 KiB
Python
# ***************************************************************
|
|
# Copyright (c) 2021 Jittor. All Rights Reserved.
|
|
# Maintainers:
|
|
# Guowei Yang <471184555@qq.com>
|
|
# Dun Liang <randonlang@gmail.com>.
|
|
#
|
|
# This file is subject to the terms and conditions defined in
|
|
# file 'LICENSE.txt', which is part of this source code package.
|
|
# ***************************************************************
|
|
import jittor as jt
|
|
import numpy as np
|
|
import math
|
|
|
|
def eye(shape, dtype):
|
|
return jt.array(np.identity(shape[0])).unary(dtype)
|
|
|
|
def eye_(var):
|
|
return var.assign(eye(var.shape, var.dtype))
|
|
|
|
def constant(shape, dtype, value=0.0):
|
|
return jt.array(value).unary(dtype).broadcast(shape)
|
|
|
|
def constant_(var, value=0.0):
|
|
return var.assign(constant(var.shape, var.dtype, value))
|
|
|
|
def uniform(shape, dtype, low, high):
|
|
return jt.random(shape, dtype) * (low - high) + high
|
|
|
|
def uniform_(var, low, high):
|
|
return var.assign(uniform(var.shape, var.dtype, low, high))
|
|
|
|
def gauss(shape, dtype, mean=0.0, std=1.0):
|
|
return jt.random(shape, dtype, "normal") * std + mean
|
|
|
|
def gauss_(var, mean=0.0, std=1.0):
|
|
return var.assign(gauss(var.shape, var.dtype, mean, std))
|
|
|
|
def invariant_uniform(shape, dtype, mode="fan_in"):
|
|
assert len(shape)>1
|
|
assert mode=="fan_in" or mode=="fan_out"
|
|
|
|
matsize=1
|
|
for i in shape[2:]:
|
|
matsize *= i
|
|
fan = (shape[1] * matsize) if mode=="fan_in" else (shape[0] * matsize)
|
|
bound = math.sqrt(1.0/fan)
|
|
return uniform(shape, dtype, -bound, bound)
|
|
|
|
def invariant_uniform_(var, mode="fan_in"):
|
|
var.assign(invariant_uniform(tuple(var.shape), var.dtype, mode))
|
|
|
|
def relu_invariant_gauss(shape, dtype, mode="fan_in"):
|
|
assert len(shape)>1
|
|
assert mode=="fan_in" or mode=="fan_out"
|
|
|
|
matsize=1
|
|
for i in shape[2:]:
|
|
matsize *= i
|
|
fan = (shape[1] * matsize) if mode=="fan_in" else (shape[0] * matsize)
|
|
std = math.sqrt(2.0/fan)
|
|
return gauss(shape, dtype, 0, std)
|
|
|
|
def relu_invariant_gauss_(var, mode="fan_in"):
|
|
return var.assign(relu_invariant_gauss(tuple(var.shape), var.dtype, mode))
|
|
|
|
def calculate_std(var,mode,nonlinearity,param=0.01):
|
|
mode = mode.lower()
|
|
assert isinstance(param,(int,float))
|
|
assert var.ndim>=2
|
|
assert mode in ['fan_in', 'fan_out']
|
|
|
|
fan = var.shape[1] if mode == 'fan_in' else var.shape[0]
|
|
fan *= var[0][0].numel()
|
|
|
|
gains = {
|
|
'linear':1,
|
|
'conv1d':1,
|
|
'conv2d':1,
|
|
'conv3d':1,
|
|
'conv_transpose1d':1,
|
|
'conv_transpose2d':1,
|
|
'conv_transpose3d':1,
|
|
'sigmoid':1,
|
|
'tanh':5.0/3,
|
|
'relu':math.sqrt(2.0),
|
|
'leaky_relu':math.sqrt(2.0 / (1 + param ** 2)),
|
|
}
|
|
gain = gains[nonlinearity]
|
|
std = gain/math.sqrt(fan)
|
|
return std
|
|
|
|
|
|
def kaiming_uniform_(var, a=0, mode='fan_in', nonlinearity='leaky_relu'):
|
|
std = calculate_std(var,mode,nonlinearity,a)
|
|
bound = math.sqrt(3.0) * std
|
|
return uniform_(var,-bound, bound)
|
|
|
|
def kaiming_normal_(var, a=0, mode='fan_in', nonlinearity='leaky_relu'):
|
|
std = calculate_std(var,mode,nonlinearity,a)
|
|
return gauss_(var,0, std)
|
|
|
|
|
|
def xavier_uniform(shape, dtype, gain=1.0):
|
|
assert len(shape)>1
|
|
|
|
matsize=1
|
|
for i in shape[2:]:
|
|
matsize *= i
|
|
fan = (shape[1] * matsize) + (shape[0] * matsize)
|
|
bound = gain * math.sqrt(6.0/fan)
|
|
return uniform(shape, dtype, -bound, bound)
|
|
|
|
def xavier_uniform_(var, gain=1.0):
|
|
return var.assign(xavier_uniform(tuple(var.shape), var.dtype, gain))
|
|
|
|
def xavier_gauss(shape, dtype, gain=1.0):
|
|
assert len(shape)>1
|
|
|
|
matsize=1
|
|
for i in shape[2:]:
|
|
matsize *= i
|
|
fan = (shape[1] * matsize) + (shape[0] * matsize)
|
|
std = gain * math.sqrt(2.0/fan)
|
|
return gauss(shape, dtype, 0, std)
|
|
|
|
def xavier_gauss_(var, gain=1.0):
|
|
return var.assign(xavier_gauss(tuple(var.shape), var.dtype, gain))
|