Spaces:
Paused
Paused
| import pdb | |
| from os import path | |
| import torch | |
| 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") | |
| _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"]) | |
| # Activation names | |
| ACT_RELU = "relu" | |
| ACT_LEAKY_RELU = "leaky_relu" | |
| 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): | |
| 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 | |
| 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) | |
| 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) | |
| ctx.mark_non_differentiable(running_mean, running_var) | |
| return x, running_mean, running_var | |
| def backward(ctx, dz, _drunning_mean, _drunning_var): | |
| z, var, weight, bias = ctx.saved_tensors | |
| dz = dz.contiguous() | |
| # Undo activation | |
| _act_backward(ctx, z, dz) | |
| if ctx.training: | |
| edz, eydz = _backend.edz_eydz(z, dz, weight, bias, ctx.affine, ctx.eps) | |
| 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 | |
| class InPlaceABNSync(autograd.Function): | |
| 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 | |
| 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 | |
| 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) | |
| 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 | |
| else: | |
| dist.all_reduce(batch_size, dist.ReduceOp.SUM) | |
| ctx.factor = x.shape[0] / float(batch_size.item()) | |
| mean_all = mean.clone() * ctx.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 = 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) | |
| 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) | |
| ctx.mark_non_differentiable(running_mean, running_var) | |
| return x, running_mean, running_var | |
| def backward(ctx, dz, _drunning_mean, _drunning_var): | |
| z, var, weight, bias = ctx.saved_tensors | |
| dz = dz.contiguous() | |
| # Undo activation | |
| _act_backward(ctx, z, dz) | |
| if ctx.training: | |
| edz, eydz = _backend.edz_eydz(z, dz, weight, bias, ctx.affine, ctx.eps) | |
| edz_local = edz.clone() | |
| eydz_local = eydz.clone() | |
| if ctx.world_size > 1: | |
| edz *= ctx.factor | |
| dist.all_reduce(edz, dist.ReduceOp.SUM) | |
| eydz *= ctx.factor | |
| dist.all_reduce(eydz, dist.ReduceOp.SUM) | |
| else: | |
| edz_local = edz = dz.new_zeros(dz.size(1)) | |
| eydz_local = 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 | |
| return dx, dweight, dbias, None, None, None, None, None, None, None | |
| inplace_abn = InPlaceABN.apply | |
| inplace_abn_sync = InPlaceABNSync.apply | |
| __all__ = ["inplace_abn", "inplace_abn_sync", "ACT_RELU", "ACT_LEAKY_RELU", "ACT_ELU", "ACT_NONE"] | |