diff --git a/schp/modules/functions.py b/schp/modules/functions.py index 4faad34..7983730 100644 --- a/schp/modules/functions.py +++ b/schp/modules/functions.py @@ -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