Arulkumar03/Wheat_HEAD_Detection_Counting_ComputerVision_Model

1# Copyright (c) Facebook, Inc. and its affiliates.2import torch3from torch import nn4from torch.nn import functional as F5from torch.nn.modules.utils import _pair6 7from detectron2.layers.wrappers import _NewEmptyTensorOp8 9 10class TridentConv(nn.Module):11    def __init__(12        self,13        in_channels,14        out_channels,15        kernel_size,16        stride=1,17        paddings=0,18        dilations=1,19        groups=1,20        num_branch=1,21        test_branch_idx=-1,22        bias=False,23        norm=None,24        activation=None,25    ):26        super(TridentConv, self).__init__()27        self.in_channels = in_channels28        self.out_channels = out_channels29        self.kernel_size = _pair(kernel_size)30        self.num_branch = num_branch31        self.stride = _pair(stride)32        self.groups = groups33        self.with_bias = bias34        if isinstance(paddings, int):35            paddings = [paddings] * self.num_branch36        if isinstance(dilations, int):37            dilations = [dilations] * self.num_branch38        self.paddings = [_pair(padding) for padding in paddings]39        self.dilations = [_pair(dilation) for dilation in dilations]40        self.test_branch_idx = test_branch_idx41        self.norm = norm42        self.activation = activation43 44        assert len({self.num_branch, len(self.paddings), len(self.dilations)}) == 145 46        self.weight = nn.Parameter(47            torch.Tensor(out_channels, in_channels // groups, *self.kernel_size)48        )49        if bias:50            self.bias = nn.Parameter(torch.Tensor(out_channels))51        else:52            self.bias = None53 54        nn.init.kaiming_uniform_(self.weight, nonlinearity="relu")55        if self.bias is not None:56            nn.init.constant_(self.bias, 0)57 58    def forward(self, inputs):59        num_branch = self.num_branch if self.training or self.test_branch_idx == -1 else 160        assert len(inputs) == num_branch61 62        if inputs[0].numel() == 0:63            output_shape = [64                (i + 2 * p - (di * (k - 1) + 1)) // s + 165                for i, p, di, k, s in zip(66                    inputs[0].shape[-2:], self.padding, self.dilation, self.kernel_size, self.stride67                )68            ]69            output_shape = [input[0].shape[0], self.weight.shape[0]] + output_shape70            return [_NewEmptyTensorOp.apply(input, output_shape) for input in inputs]71 72        if self.training or self.test_branch_idx == -1:73            outputs = [74                F.conv2d(input, self.weight, self.bias, self.stride, padding, dilation, self.groups)75                for input, dilation, padding in zip(inputs, self.dilations, self.paddings)76            ]77        else:78            outputs = [79                F.conv2d(80                    inputs[0],81                    self.weight,82                    self.bias,83                    self.stride,84                    self.paddings[self.test_branch_idx],85                    self.dilations[self.test_branch_idx],86                    self.groups,87                )88            ]89 90        if self.norm is not None:91            outputs = [self.norm(x) for x in outputs]92        if self.activation is not None:93            outputs = [self.activation(x) for x in outputs]94        return outputs95 96    def extra_repr(self):97        tmpstr = "in_channels=" + str(self.in_channels)98        tmpstr += ", out_channels=" + str(self.out_channels)99        tmpstr += ", kernel_size=" + str(self.kernel_size)100        tmpstr += ", num_branch=" + str(self.num_branch)101        tmpstr += ", test_branch_idx=" + str(self.test_branch_idx)102        tmpstr += ", stride=" + str(self.stride)103        tmpstr += ", paddings=" + str(self.paddings)104        tmpstr += ", dilations=" + str(self.dilations)105        tmpstr += ", groups=" + str(self.groups)106        tmpstr += ", bias=" + str(self.with_bias)107        return tmpstr108