Focus width and height information into channel space.
| 186 | |
| 187 | |
| 188 | class Focus(nn.Module): |
| 189 | """Focus width and height information into channel space.""" |
| 190 | |
| 191 | def __init__(self, in_channels, out_channels, ksize=1, stride=1, act="silu"): |
| 192 | super().__init__() |
| 193 | self.conv = BaseConv(in_channels * 4, out_channels, ksize, stride, act=act) |
| 194 | |
| 195 | def forward(self, x): |
| 196 | # shape of x (b,c,w,h) -> y(b,4c,w/2,h/2) |
| 197 | patch_top_left = x[..., ::2, ::2] |
| 198 | patch_top_right = x[..., ::2, 1::2] |
| 199 | patch_bot_left = x[..., 1::2, ::2] |
| 200 | patch_bot_right = x[..., 1::2, 1::2] |
| 201 | x = torch.cat( |
| 202 | ( |
| 203 | patch_top_left, |
| 204 | patch_bot_left, |
| 205 | patch_top_right, |
| 206 | patch_bot_right, |
| 207 | ), |
| 208 | dim=1, |
| 209 | ) |
| 210 | return self.conv(x) |