Arulkumar03/Fox_Sheep_Detector_Computer_Vision_model
0
1# Copyright (c) Facebook, Inc. and its affiliates.2import math3from functools import lru_cache4import torch5from torch import nn6from torch.autograd import Function7from torch.autograd.function import once_differentiable8from torch.nn.modules.utils import _pair9from torchvision.ops import deform_conv2d10 11from detectron2.utils.develop import create_dummy_class, create_dummy_func12 13from .wrappers import _NewEmptyTensorOp14 15 16class _DeformConv(Function):17 @staticmethod18 def forward(19 ctx,20 input,21 offset,22 weight,23 stride=1,24 padding=0,25 dilation=1,26 groups=1,27 deformable_groups=1,28 im2col_step=64,29 ):30 if input is not None and input.dim() != 4:31 raise ValueError(32 "Expected 4D tensor as input, got {}D tensor instead.".format(input.dim())33 )34 ctx.stride = _pair(stride)35 ctx.padding = _pair(padding)36 ctx.dilation = _pair(dilation)37 ctx.groups = groups38 ctx.deformable_groups = deformable_groups39 ctx.im2col_step = im2col_step40 41 ctx.save_for_backward(input, offset, weight)42 43 output = input.new_empty(44 _DeformConv._output_size(input, weight, ctx.padding, ctx.dilation, ctx.stride)45 )46 47 ctx.bufs_ = [input.new_empty(0), input.new_empty(0)] # columns, ones48 49 if not input.is_cuda:50 # TODO: let torchvision support full features of our deformconv.51 if deformable_groups != 1:52 raise NotImplementedError(53 "Deformable Conv with deformable_groups != 1 is not supported on CPUs!"54 )55 return deform_conv2d(56 input, offset, weight, stride=stride, padding=padding, dilation=dilation57 )58 else:59 cur_im2col_step = _DeformConv._cal_im2col_step(input.shape[0], ctx.im2col_step)60 assert (input.shape[0] % cur_im2col_step) == 0, "im2col step must divide batchsize"61 62 _C.deform_conv_forward(63 input,64 weight,65 offset,66 output,67 ctx.bufs_[0],68 ctx.bufs_[1],69 weight.size(3),70 weight.size(2),71 ctx.stride[1],72 ctx.stride[0],73 ctx.padding[1],74 ctx.padding[0],75 ctx.dilation[1],76 ctx.dilation[0],77 ctx.groups,78 ctx.deformable_groups,79 cur_im2col_step,80 )81 return output82 83 @staticmethod84 @once_differentiable85 def backward(ctx, grad_output):86 input, offset, weight = ctx.saved_tensors87 88 grad_input = grad_offset = grad_weight = None89 90 if not grad_output.is_cuda:91 raise NotImplementedError("Deformable Conv is not supported on CPUs!")92 else:93 cur_im2col_step = _DeformConv._cal_im2col_step(input.shape[0], ctx.im2col_step)94 assert (input.shape[0] % cur_im2col_step) == 0, "im2col step must divide batchsize"95 96 if ctx.needs_input_grad[0] or ctx.needs_input_grad[1]:97 grad_input = torch.zeros_like(input)98 grad_offset = torch.zeros_like(offset)99 _C.deform_conv_backward_input(100 input,101 offset,102 grad_output,103 grad_input,104 grad_offset,105 weight,106 ctx.bufs_[0],107 weight.size(3),108 weight.size(2),109 ctx.stride[1],110 ctx.stride[0],111 ctx.padding[1],112 ctx.padding[0],113 ctx.dilation[1],114 ctx.dilation[0],115 ctx.groups,116 ctx.deformable_groups,117 cur_im2col_step,118 )119 120 if ctx.needs_input_grad[2]:121 grad_weight = torch.zeros_like(weight)122 _C.deform_conv_backward_filter(123 input,124 offset,125 grad_output,126 grad_weight,127 ctx.bufs_[0],128 ctx.bufs_[1],129 weight.size(3),130 weight.size(2),131 ctx.stride[1],132 ctx.stride[0],133 ctx.padding[1],134 ctx.padding[0],135 ctx.dilation[1],136 ctx.dilation[0],137 ctx.groups,138 ctx.deformable_groups,139 1,140 cur_im2col_step,141 )142 143 return grad_input, grad_offset, grad_weight, None, None, None, None, None, None144 145 @staticmethod146 def _output_size(input, weight, padding, dilation, stride):147 channels = weight.size(0)148 output_size = (input.size(0), channels)149 for d in range(input.dim() - 2):150 in_size = input.size(d + 2)151 pad = padding[d]152 kernel = dilation[d] * (weight.size(d + 2) - 1) + 1153 stride_ = stride[d]154 output_size += ((in_size + (2 * pad) - kernel) // stride_ + 1,)155 if not all(map(lambda s: s > 0, output_size)):156 raise ValueError(157 "convolution input is too small (output would be {})".format(158 "x".join(map(str, output_size))159 )160 )161 return output_size162 163 @staticmethod164 @lru_cache(maxsize=128)165 def _cal_im2col_step(input_size, default_size):166 """167 Calculate proper im2col step size, which should be divisible by input_size and not larger168 than prefer_size. Meanwhile the step size should be as large as possible to be more169 efficient. So we choose the largest one among all divisors of input_size which are smaller170 than prefer_size.171 :param input_size: input batch size .172 :param default_size: default preferred im2col step size.173 :return: the largest proper step size.174 """175 if input_size <= default_size:176 return input_size177 best_step = 1178 for step in range(2, min(int(math.sqrt(input_size)) + 1, default_size)):179 if input_size % step == 0:180 if input_size // step <= default_size:181 return input_size // step182 best_step = step183 184 return best_step185 186 187class _ModulatedDeformConv(Function):188 @staticmethod189 def forward(190 ctx,191 input,192 offset,193 mask,194 weight,195 bias=None,196 stride=1,197 padding=0,198 dilation=1,199 groups=1,200 deformable_groups=1,201 ):202 ctx.stride = stride203 ctx.padding = padding204 ctx.dilation = dilation205 ctx.groups = groups206 ctx.deformable_groups = deformable_groups207 ctx.with_bias = bias is not None208 if not ctx.with_bias:209 bias = input.new_empty(1) # fake tensor210 if not input.is_cuda:211 raise NotImplementedError("Deformable Conv is not supported on CPUs!")212 if (213 weight.requires_grad214 or mask.requires_grad215 or offset.requires_grad216 or input.requires_grad217 ):218 ctx.save_for_backward(input, offset, mask, weight, bias)219 output = input.new_empty(_ModulatedDeformConv._infer_shape(ctx, input, weight))220 ctx._bufs = [input.new_empty(0), input.new_empty(0)]221 _C.modulated_deform_conv_forward(222 input,223 weight,224 bias,225 ctx._bufs[0],226 offset,227 mask,228 output,229 ctx._bufs[1],230 weight.shape[2],231 weight.shape[3],232 ctx.stride,233 ctx.stride,234 ctx.padding,235 ctx.padding,236 ctx.dilation,237 ctx.dilation,238 ctx.groups,239 ctx.deformable_groups,240 ctx.with_bias,241 )242 return output243 244 @staticmethod245 @once_differentiable246 def backward(ctx, grad_output):247 if not grad_output.is_cuda:248 raise NotImplementedError("Deformable Conv is not supported on CPUs!")249 input, offset, mask, weight, bias = ctx.saved_tensors250 grad_input = torch.zeros_like(input)251 grad_offset = torch.zeros_like(offset)252 grad_mask = torch.zeros_like(mask)253 grad_weight = torch.zeros_like(weight)254 grad_bias = torch.zeros_like(bias)255 _C.modulated_deform_conv_backward(256 input,257 weight,258 bias,259 ctx._bufs[0],260 offset,261 mask,262 ctx._bufs[1],263 grad_input,264 grad_weight,265 grad_bias,266 grad_offset,267 grad_mask,268 grad_output,269 weight.shape[2],270 weight.shape[3],271 ctx.stride,272 ctx.stride,273 ctx.padding,274 ctx.padding,275 ctx.dilation,276 ctx.dilation,277 ctx.groups,278 ctx.deformable_groups,279 ctx.with_bias,280 )281 if not ctx.with_bias:282 grad_bias = None283 284 return (285 grad_input,286 grad_offset,287 grad_mask,288 grad_weight,289 grad_bias,290 None,291 None,292 None,293 None,294 None,295 )296 297 @staticmethod298 def _infer_shape(ctx, input, weight):299 n = input.size(0)300 channels_out = weight.size(0)301 height, width = input.shape[2:4]302 kernel_h, kernel_w = weight.shape[2:4]303 height_out = (304 height + 2 * ctx.padding - (ctx.dilation * (kernel_h - 1) + 1)305 ) // ctx.stride + 1306 width_out = (307 width + 2 * ctx.padding - (ctx.dilation * (kernel_w - 1) + 1)308 ) // ctx.stride + 1309 return n, channels_out, height_out, width_out310 311 312deform_conv = _DeformConv.apply313modulated_deform_conv = _ModulatedDeformConv.apply314 315 316class DeformConv(nn.Module):317 def __init__(318 self,319 in_channels,320 out_channels,321 kernel_size,322 stride=1,323 padding=0,324 dilation=1,325 groups=1,326 deformable_groups=1,327 bias=False,328 norm=None,329 activation=None,330 ):331 """332 Deformable convolution from :paper:`deformconv`.333 334 Arguments are similar to :class:`Conv2D`. Extra arguments:335 336 Args:337 deformable_groups (int): number of groups used in deformable convolution.338 norm (nn.Module, optional): a normalization layer339 activation (callable(Tensor) -> Tensor): a callable activation function340 """341 super(DeformConv, self).__init__()342 343 assert not bias344 assert in_channels % groups == 0, "in_channels {} cannot be divisible by groups {}".format(345 in_channels, groups346 )347 assert (348 out_channels % groups == 0349 ), "out_channels {} cannot be divisible by groups {}".format(out_channels, groups)350 351 self.in_channels = in_channels352 self.out_channels = out_channels353 self.kernel_size = _pair(kernel_size)354 self.stride = _pair(stride)355 self.padding = _pair(padding)356 self.dilation = _pair(dilation)357 self.groups = groups358 self.deformable_groups = deformable_groups359 self.norm = norm360 self.activation = activation361 362 self.weight = nn.Parameter(363 torch.Tensor(out_channels, in_channels // self.groups, *self.kernel_size)364 )365 self.bias = None366 367 nn.init.kaiming_uniform_(self.weight, nonlinearity="relu")368 369 def forward(self, x, offset):370 if x.numel() == 0:371 # When input is empty, we want to return a empty tensor with "correct" shape,372 # So that the following operations will not panic373 # if they check for the shape of the tensor.374 # This computes the height and width of the output tensor375 output_shape = [376 (i + 2 * p - (di * (k - 1) + 1)) // s + 1377 for i, p, di, k, s in zip(378 x.shape[-2:], self.padding, self.dilation, self.kernel_size, self.stride379 )380 ]381 output_shape = [x.shape[0], self.weight.shape[0]] + output_shape382 return _NewEmptyTensorOp.apply(x, output_shape)383 384 x = deform_conv(385 x,386 offset,387 self.weight,388 self.stride,389 self.padding,390 self.dilation,391 self.groups,392 self.deformable_groups,393 )394 if self.norm is not None:395 x = self.norm(x)396 if self.activation is not None:397 x = self.activation(x)398 return x399 400 def extra_repr(self):401 tmpstr = "in_channels=" + str(self.in_channels)402 tmpstr += ", out_channels=" + str(self.out_channels)403 tmpstr += ", kernel_size=" + str(self.kernel_size)404 tmpstr += ", stride=" + str(self.stride)405 tmpstr += ", padding=" + str(self.padding)406 tmpstr += ", dilation=" + str(self.dilation)407 tmpstr += ", groups=" + str(self.groups)408 tmpstr += ", deformable_groups=" + str(self.deformable_groups)409 tmpstr += ", bias=False"410 return tmpstr411 412 413class ModulatedDeformConv(nn.Module):414 def __init__(415 self,416 in_channels,417 out_channels,418 kernel_size,419 stride=1,420 padding=0,421 dilation=1,422 groups=1,423 deformable_groups=1,424 bias=True,425 norm=None,426 activation=None,427 ):428 """429 Modulated deformable convolution from :paper:`deformconv2`.430 431 Arguments are similar to :class:`Conv2D`. Extra arguments:432 433 Args:434 deformable_groups (int): number of groups used in deformable convolution.435 norm (nn.Module, optional): a normalization layer436 activation (callable(Tensor) -> Tensor): a callable activation function437 """438 super(ModulatedDeformConv, self).__init__()439 self.in_channels = in_channels440 self.out_channels = out_channels441 self.kernel_size = _pair(kernel_size)442 self.stride = stride443 self.padding = padding444 self.dilation = dilation445 self.groups = groups446 self.deformable_groups = deformable_groups447 self.with_bias = bias448 self.norm = norm449 self.activation = activation450 451 self.weight = nn.Parameter(452 torch.Tensor(out_channels, in_channels // groups, *self.kernel_size)453 )454 if bias:455 self.bias = nn.Parameter(torch.Tensor(out_channels))456 else:457 self.bias = None458 459 nn.init.kaiming_uniform_(self.weight, nonlinearity="relu")460 if self.bias is not None:461 nn.init.constant_(self.bias, 0)462 463 def forward(self, x, offset, mask):464 if x.numel() == 0:465 output_shape = [466 (i + 2 * p - (di * (k - 1) + 1)) // s + 1467 for i, p, di, k, s in zip(468 x.shape[-2:], self.padding, self.dilation, self.kernel_size, self.stride469 )470 ]471 output_shape = [x.shape[0], self.weight.shape[0]] + output_shape472 return _NewEmptyTensorOp.apply(x, output_shape)473 474 x = modulated_deform_conv(475 x,476 offset,477 mask,478 self.weight,479 self.bias,480 self.stride,481 self.padding,482 self.dilation,483 self.groups,484 self.deformable_groups,485 )486 if self.norm is not None:487 x = self.norm(x)488 if self.activation is not None:489 x = self.activation(x)490 return x491 492 def extra_repr(self):493 tmpstr = "in_channels=" + str(self.in_channels)494 tmpstr += ", out_channels=" + str(self.out_channels)495 tmpstr += ", kernel_size=" + str(self.kernel_size)496 tmpstr += ", stride=" + str(self.stride)497 tmpstr += ", padding=" + str(self.padding)498 tmpstr += ", dilation=" + str(self.dilation)499 tmpstr += ", groups=" + str(self.groups)500 tmpstr += ", deformable_groups=" + str(self.deformable_groups)501 tmpstr += ", bias=" + str(self.with_bias)502 return tmpstr503 504 505try:506 from detectron2 import _C507except ImportError:508 # TODO: register ops natively so there is no need to import _C.509 _msg = "detectron2 is not compiled successfully, please build following the instructions!"510 _args = ("detectron2._C", _msg)511 DeformConv = create_dummy_class("DeformConv", *_args)512 ModulatedDeformConv = create_dummy_class("ModulatedDeformConv", *_args)513 deform_conv = create_dummy_func("deform_conv", *_args)514 modulated_deform_conv = create_dummy_func("modulated_deform_conv", *_args)515