Update functions.py
This commit is contained in:
Jinwoo Park (Curt)
+99
-171
@@ -4,27 +4,7 @@ import torch.distributed as dist
|
||||
import torch.autograd as autograd
|
||||
import torch.cuda.comm as comm
|
||||
from torch.autograd.function import once_differentiable
|
||||
from torch.utils.cpp_extension import load
|
||||
|
||||
_src_path = path.join(path.dirname(path.abspath(__file__)), "src")
|
||||
|
||||
if torch.cuda.is_available():
|
||||
_backend = load(name="inplace_abn",
|
||||
extra_cflags=["-O3"],
|
||||
sources=[path.join(_src_path, f) for f in [
|
||||
"inplace_abn.cpp",
|
||||
"inplace_abn_cpu.cpp",
|
||||
"inplace_abn_cuda.cu",
|
||||
"inplace_abn_cuda_half.cu"
|
||||
]],
|
||||
extra_cuda_cflags=["--expt-extended-lambda"])
|
||||
else:
|
||||
_backend = load(name="inplace_abn",
|
||||
extra_cflags=["-O3"],
|
||||
sources=[path.join(_src_path, f) for f in [
|
||||
"inplace_abn_cpu_only.cpp",
|
||||
"inplace_abn_cpu.cpp"
|
||||
]])
|
||||
|
||||
# Activation names
|
||||
ACT_RELU = "relu"
|
||||
@@ -33,128 +13,72 @@ ACT_ELU = "elu"
|
||||
ACT_NONE = "none"
|
||||
|
||||
|
||||
def _check(fn, *args, **kwargs):
|
||||
success = fn(*args, **kwargs)
|
||||
if not success:
|
||||
raise RuntimeError("CUDA Error encountered in {}".format(fn))
|
||||
|
||||
|
||||
def _broadcast_shape(x):
|
||||
out_size = []
|
||||
for i, s in enumerate(x.size()):
|
||||
if i != 1:
|
||||
out_size.append(1)
|
||||
else:
|
||||
out_size.append(s)
|
||||
return out_size
|
||||
|
||||
|
||||
def _reduce(x):
|
||||
if len(x.size()) == 2:
|
||||
return x.sum(dim=0)
|
||||
else:
|
||||
n, c = x.size()[0:2]
|
||||
return x.contiguous().view((n, c, -1)).sum(2).sum(0)
|
||||
|
||||
|
||||
def _count_samples(x):
|
||||
count = 1
|
||||
for i, s in enumerate(x.size()):
|
||||
if i != 1:
|
||||
count *= s
|
||||
return count
|
||||
|
||||
|
||||
def _act_forward(ctx, x):
|
||||
if ctx.activation == ACT_LEAKY_RELU:
|
||||
_backend.leaky_relu_forward(x, ctx.slope)
|
||||
elif ctx.activation == ACT_ELU:
|
||||
_backend.elu_forward(x)
|
||||
elif ctx.activation == ACT_NONE:
|
||||
pass
|
||||
|
||||
|
||||
def _act_backward(ctx, x, dx):
|
||||
if ctx.activation == ACT_LEAKY_RELU:
|
||||
_backend.leaky_relu_backward(x, dx, ctx.slope)
|
||||
elif ctx.activation == ACT_ELU:
|
||||
_backend.elu_backward(x, dx)
|
||||
elif ctx.activation == ACT_NONE:
|
||||
pass
|
||||
|
||||
|
||||
class InPlaceABN(autograd.Function):
|
||||
@staticmethod
|
||||
def forward(ctx, x, weight, bias, running_mean, running_var,
|
||||
training=True, momentum=0.1, eps=1e-05, activation=ACT_LEAKY_RELU, slope=0.01):
|
||||
# Save context
|
||||
training=True, momentum=0.1, eps=1e-05, activation='leaky_relu', slope=0.01):
|
||||
ctx.training = training
|
||||
ctx.momentum = momentum
|
||||
ctx.eps = eps
|
||||
ctx.activation = activation
|
||||
ctx.slope = slope
|
||||
ctx.affine = weight is not None and bias is not None
|
||||
|
||||
# Prepare inputs
|
||||
count = _count_samples(x)
|
||||
x = x.contiguous()
|
||||
weight = weight.contiguous() if ctx.affine else x.new_empty(0)
|
||||
bias = bias.contiguous() if ctx.affine else x.new_empty(0)
|
||||
|
||||
if ctx.training:
|
||||
mean, var = _backend.mean_var(x)
|
||||
|
||||
# Update running stats
|
||||
running_mean.mul_((1 - ctx.momentum)).add_(ctx.momentum * mean)
|
||||
running_var.mul_((1 - ctx.momentum)).add_(ctx.momentum * var * count / (count - 1))
|
||||
|
||||
# Mark in-place modified tensors
|
||||
ctx.mark_dirty(x, running_mean, running_var)
|
||||
|
||||
if training:
|
||||
mean = x.mean(dim=[0, 2, 3], keepdim=True)
|
||||
var = x.var(dim=[0, 2, 3], keepdim=True)
|
||||
|
||||
running_mean = (1 - momentum) * running_mean + momentum * mean
|
||||
running_var = (1 - momentum) * running_var + momentum * var * x.size(0) / (x.size(0) - 1)
|
||||
else:
|
||||
mean, var = running_mean.contiguous(), running_var.contiguous()
|
||||
ctx.mark_dirty(x)
|
||||
|
||||
# BN forward + activation
|
||||
_backend.forward(x, mean, var, weight, bias, ctx.affine, ctx.eps)
|
||||
_act_forward(ctx, x)
|
||||
|
||||
# Output
|
||||
ctx.var = var
|
||||
ctx.save_for_backward(x, var, weight, bias)
|
||||
mean = running_mean
|
||||
var = running_var
|
||||
|
||||
x_norm = (x - mean.view(1, -1, 1, 1)) / torch.sqrt(var.view(1, -1, 1, 1) + eps)
|
||||
|
||||
if ctx.affine:
|
||||
x_norm = x_norm * weight.view(1, -1, 1, 1) + bias.view(1, -1, 1, 1)
|
||||
|
||||
# Apply activation
|
||||
if ctx.activation == 'leaky_relu':
|
||||
x_norm = torch.where(x_norm > 0, x_norm, x_norm * ctx.slope)
|
||||
|
||||
ctx.save_for_backward(x_norm, var, weight, bias)
|
||||
ctx.mark_non_differentiable(running_mean, running_var)
|
||||
return x, running_mean, running_var
|
||||
|
||||
|
||||
return x_norm
|
||||
|
||||
@staticmethod
|
||||
@once_differentiable
|
||||
def backward(ctx, dz, _drunning_mean, _drunning_var):
|
||||
def backward(ctx, dz):
|
||||
z, var, weight, bias = ctx.saved_tensors
|
||||
dz = dz.contiguous()
|
||||
|
||||
|
||||
# Undo activation
|
||||
_act_backward(ctx, z, dz)
|
||||
|
||||
if ctx.activation == 'leaky_relu':
|
||||
dz = torch.where(z > 0, dz, dz * ctx.slope)
|
||||
|
||||
# Compute gradients
|
||||
if ctx.training:
|
||||
edz, eydz = _backend.edz_eydz(z, dz, weight, bias, ctx.affine, ctx.eps)
|
||||
edz, eydz = dz.mean(dim=[0, 2, 3]), (dz * z).mean(dim=[0, 2, 3])
|
||||
else:
|
||||
# TODO: implement simplified CUDA backward for inference mode
|
||||
edz = dz.new_zeros(dz.size(1))
|
||||
eydz = dz.new_zeros(dz.size(1))
|
||||
|
||||
dx = _backend.backward(z, dz, var, weight, bias, edz, eydz, ctx.affine, ctx.eps)
|
||||
# dweight = eydz * weight.sign() if ctx.affine else None
|
||||
dweight = eydz if ctx.affine else None
|
||||
if dweight is not None:
|
||||
dweight[weight < 0] *= -1
|
||||
dbias = edz if ctx.affine else None
|
||||
|
||||
return dx, dweight, dbias, None, None, None, None, None, None, None
|
||||
edz, eydz = dz.new_zeros(dz.size(1)), dz.new_zeros(dz.size(1))
|
||||
|
||||
dx = dz * (1 / torch.sqrt(var + ctx.eps))
|
||||
|
||||
if ctx.affine:
|
||||
dweight = eydz if ctx.affine else None
|
||||
dbias = edz if ctx.affine else None
|
||||
else:
|
||||
dweight = None
|
||||
dbias = None
|
||||
|
||||
return dx, dweight, dbias, None, None, None, None, None, None
|
||||
|
||||
|
||||
class InPlaceABNSync(autograd.Function):
|
||||
@classmethod
|
||||
def forward(cls, ctx, x, weight, bias, running_mean, running_var,
|
||||
training=True, momentum=0.1, eps=1e-05, activation=ACT_LEAKY_RELU, slope=0.01, equal_batches=True):
|
||||
# Save context
|
||||
training=True, momentum=0.1, eps=1e-05, activation='leaky_relu', slope=0.01, equal_batches=True):
|
||||
ctx.training = training
|
||||
ctx.momentum = momentum
|
||||
ctx.eps = eps
|
||||
@@ -162,89 +86,93 @@ class InPlaceABNSync(autograd.Function):
|
||||
ctx.slope = slope
|
||||
ctx.affine = weight is not None and bias is not None
|
||||
|
||||
# Prepare inputs
|
||||
ctx.world_size = dist.get_world_size() if dist.is_initialized() else 1
|
||||
|
||||
# count = _count_samples(x)
|
||||
batch_size = x.new_tensor([x.shape[0]], dtype=torch.long)
|
||||
if training:
|
||||
mean = x.mean(dim=[0, 2, 3], keepdim=True)
|
||||
var = x.var(dim=[0, 2, 3], keepdim=True)
|
||||
|
||||
x = x.contiguous()
|
||||
weight = weight.contiguous() if ctx.affine else x.new_empty(0)
|
||||
bias = bias.contiguous() if ctx.affine else x.new_empty(0)
|
||||
|
||||
if ctx.training:
|
||||
mean, var = _backend.mean_var(x)
|
||||
if ctx.world_size > 1:
|
||||
# get global batch size
|
||||
if equal_batches:
|
||||
batch_size *= ctx.world_size
|
||||
batch_size = x.size(0) * ctx.world_size
|
||||
else:
|
||||
batch_size = x.new_tensor([x.size(0)], dtype=torch.long)
|
||||
dist.all_reduce(batch_size, dist.ReduceOp.SUM)
|
||||
batch_size = batch_size.item()
|
||||
|
||||
ctx.factor = x.shape[0] / float(batch_size.item())
|
||||
factor = x.size(0) / batch_size
|
||||
|
||||
mean_all = mean.clone() * ctx.factor
|
||||
mean_all = mean.clone() * factor
|
||||
dist.all_reduce(mean_all, dist.ReduceOp.SUM)
|
||||
|
||||
var_all = (var + (mean - mean_all) ** 2) * ctx.factor
|
||||
dist.all_reduce(var_all, dist.ReduceOp.SUM)
|
||||
|
||||
mean = mean_all
|
||||
|
||||
var_all = (var + (mean - mean_all) ** 2) * factor
|
||||
dist.all_reduce(var_all, dist.ReduceOp.SUM)
|
||||
var = var_all
|
||||
|
||||
# Update running stats
|
||||
running_mean.mul_((1 - ctx.momentum)).add_(ctx.momentum * mean)
|
||||
count = batch_size.item() * x.view(x.shape[0], x.shape[1], -1).shape[-1]
|
||||
running_var.mul_((1 - ctx.momentum)).add_(ctx.momentum * var * (float(count) / (count - 1)))
|
||||
|
||||
# Mark in-place modified tensors
|
||||
ctx.mark_dirty(x, running_mean, running_var)
|
||||
running_mean = (1 - momentum) * running_mean + momentum * mean
|
||||
running_var = (1 - momentum) * running_var + momentum * var * batch_size / (batch_size - 1)
|
||||
else:
|
||||
running_mean = (1 - momentum) * running_mean + momentum * mean
|
||||
running_var = (1 - momentum) * running_var + momentum * var * x.size(0) / (x.size(0) - 1)
|
||||
else:
|
||||
mean, var = running_mean.contiguous(), running_var.contiguous()
|
||||
ctx.mark_dirty(x)
|
||||
mean = running_mean
|
||||
var = running_var
|
||||
|
||||
# BN forward + activation
|
||||
_backend.forward(x, mean, var, weight, bias, ctx.affine, ctx.eps)
|
||||
_act_forward(ctx, x)
|
||||
x_norm = (x - mean.view(1, -1, 1, 1)) / torch.sqrt(var.view(1, -1, 1, 1) + eps)
|
||||
|
||||
# Output
|
||||
ctx.var = var
|
||||
ctx.save_for_backward(x, var, weight, bias)
|
||||
if ctx.affine:
|
||||
x_norm = x_norm * weight.view(1, -1, 1, 1) + bias.view(1, -1, 1, 1)
|
||||
|
||||
# Apply activation
|
||||
if ctx.activation == 'leaky_relu':
|
||||
x_norm = torch.where(x_norm > 0, x_norm, x_norm * ctx.slope)
|
||||
|
||||
ctx.save_for_backward(x_norm, var, weight, bias)
|
||||
ctx.mark_non_differentiable(running_mean, running_var)
|
||||
return x, running_mean, running_var
|
||||
|
||||
return x_norm, running_mean, running_var
|
||||
|
||||
@staticmethod
|
||||
@once_differentiable
|
||||
def backward(ctx, dz, _drunning_mean, _drunning_var):
|
||||
def backward(ctx, dz):
|
||||
z, var, weight, bias = ctx.saved_tensors
|
||||
dz = dz.contiguous()
|
||||
|
||||
# Undo activation
|
||||
_act_backward(ctx, z, dz)
|
||||
if ctx.activation == 'leaky_relu':
|
||||
dz = torch.where(z > 0, dz, dz * ctx.slope)
|
||||
|
||||
# Compute gradients
|
||||
if ctx.training:
|
||||
edz, eydz = _backend.edz_eydz(z, dz, weight, bias, ctx.affine, ctx.eps)
|
||||
edz_local = edz.clone()
|
||||
eydz_local = eydz.clone()
|
||||
edz_local = dz.mean(dim=[0, 2, 3])
|
||||
eydz_local = (dz * z).mean(dim=[0, 2, 3])
|
||||
|
||||
if ctx.world_size > 1:
|
||||
edz *= ctx.factor
|
||||
edz = dz.new_zeros(dz.size(1))
|
||||
eydz = dz.new_zeros(dz.size(1))
|
||||
|
||||
edz_local *= ctx.world_size
|
||||
eydz_local *= ctx.world_size
|
||||
|
||||
dist.all_reduce(edz, dist.ReduceOp.SUM)
|
||||
|
||||
eydz *= ctx.factor
|
||||
dist.all_reduce(eydz, dist.ReduceOp.SUM)
|
||||
else:
|
||||
edz = edz_local
|
||||
eydz = eydz_local
|
||||
else:
|
||||
edz_local = edz = dz.new_zeros(dz.size(1))
|
||||
eydz_local = eydz = dz.new_zeros(dz.size(1))
|
||||
edz = dz.new_zeros(dz.size(1))
|
||||
eydz = dz.new_zeros(dz.size(1))
|
||||
|
||||
dx = _backend.backward(z, dz, var, weight, bias, edz, eydz, ctx.affine, ctx.eps)
|
||||
# dweight = eydz_local * weight.sign() if ctx.affine else None
|
||||
dweight = eydz_local if ctx.affine else None
|
||||
if dweight is not None:
|
||||
dweight[weight < 0] *= -1
|
||||
dbias = edz_local if ctx.affine else None
|
||||
dx = dz * (1 / torch.sqrt(var + ctx.eps))
|
||||
|
||||
return dx, dweight, dbias, None, None, None, None, None, None, None
|
||||
if ctx.affine:
|
||||
dweight = eydz_local if ctx.affine else None
|
||||
dbias = edz_local if ctx.affine else None
|
||||
else:
|
||||
dweight = None
|
||||
dbias = None
|
||||
|
||||
return dx, dweight, dbias, None, None, None, None, None, None
|
||||
|
||||
|
||||
inplace_abn = InPlaceABN.apply
|
||||
|
||||
Reference in New Issue
Block a user