| 81 | |
| 82 | |
| 83 | class DualPathBlock(nn.Module): |
| 84 | def __init__( |
| 85 | self, in_chs, num_1x1_a, num_3x3_b, num_1x1_c, inc, groups, block_type='normal', b=False): |
| 86 | super(DualPathBlock, self).__init__() |
| 87 | self.num_1x1_c = num_1x1_c |
| 88 | self.inc = inc |
| 89 | self.b = b |
| 90 | if block_type == 'proj': |
| 91 | self.key_stride = 1 |
| 92 | self.has_proj = True |
| 93 | elif block_type == 'down': |
| 94 | self.key_stride = 2 |
| 95 | self.has_proj = True |
| 96 | else: |
| 97 | assert block_type == 'normal' |
| 98 | self.key_stride = 1 |
| 99 | self.has_proj = False |
| 100 | |
| 101 | self.c1x1_w_s1 = None |
| 102 | self.c1x1_w_s2 = None |
| 103 | if self.has_proj: |
| 104 | # Using different member names here to allow easier parameter key matching for conversion |
| 105 | if self.key_stride == 2: |
| 106 | self.c1x1_w_s2 = BnActConv2d( |
| 107 | in_chs=in_chs, out_chs=num_1x1_c + 2 * inc, kernel_size=1, stride=2) |
| 108 | else: |
| 109 | self.c1x1_w_s1 = BnActConv2d( |
| 110 | in_chs=in_chs, out_chs=num_1x1_c + 2 * inc, kernel_size=1, stride=1) |
| 111 | |
| 112 | self.c1x1_a = BnActConv2d(in_chs=in_chs, out_chs=num_1x1_a, kernel_size=1, stride=1) |
| 113 | self.c3x3_b = BnActConv2d( |
| 114 | in_chs=num_1x1_a, out_chs=num_3x3_b, kernel_size=3, stride=self.key_stride, groups=groups) |
| 115 | if b: |
| 116 | self.c1x1_c = CatBnAct(in_chs=num_3x3_b) |
| 117 | self.c1x1_c1 = create_conv2d(num_3x3_b, num_1x1_c, kernel_size=1) |
| 118 | self.c1x1_c2 = create_conv2d(num_3x3_b, inc, kernel_size=1) |
| 119 | else: |
| 120 | self.c1x1_c = BnActConv2d(in_chs=num_3x3_b, out_chs=num_1x1_c + inc, kernel_size=1, stride=1) |
| 121 | self.c1x1_c1 = None |
| 122 | self.c1x1_c2 = None |
| 123 | |
| 124 | @torch.jit._overload_method # noqa: F811 |
| 125 | def forward(self, x): |
| 126 | # type: (Tuple[torch.Tensor, torch.Tensor]) -> Tuple[torch.Tensor, torch.Tensor] |
| 127 | pass |
| 128 | |
| 129 | @torch.jit._overload_method # noqa: F811 |
| 130 | def forward(self, x): |
| 131 | # type: (torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor] |
| 132 | pass |
| 133 | |
| 134 | def forward(self, x) -> Tuple[torch.Tensor, torch.Tensor]: |
| 135 | if isinstance(x, tuple): |
| 136 | x_in = torch.cat(x, dim=1) |
| 137 | else: |
| 138 | x_in = x |
| 139 | if self.c1x1_w_s1 is None and self.c1x1_w_s2 is None: |
| 140 | # self.has_proj == False, torchscript requires condition on module == None |