Modification of https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/levit.py: ConvNorm such that ConvTranspose2d is used instead of Conv2d.
| 53 | |
| 54 | |
| 55 | class ConvTransposeNorm(nn.Sequential): |
| 56 | """ |
| 57 | Modification of |
| 58 | https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/levit.py: ConvNorm |
| 59 | such that ConvTranspose2d is used instead of Conv2d. |
| 60 | """ |
| 61 | |
| 62 | def __init__( |
| 63 | self, in_chs, out_chs, kernel_size=1, stride=1, pad=0, dilation=1, |
| 64 | groups=1, bn_weight_init=1): |
| 65 | super().__init__() |
| 66 | self.add_module('c', |
| 67 | nn.ConvTranspose2d(in_chs, out_chs, kernel_size, stride, pad, dilation, groups, bias=False)) |
| 68 | self.add_module('bn', nn.BatchNorm2d(out_chs)) |
| 69 | |
| 70 | nn.init.constant_(self.bn.weight, bn_weight_init) |
| 71 | |
| 72 | @torch.no_grad() |
| 73 | def fuse(self): |
| 74 | c, bn = self._modules.values() |
| 75 | w = bn.weight / (bn.running_var + bn.eps) ** 0.5 |
| 76 | w = c.weight * w[:, None, None, None] |
| 77 | b = bn.bias - bn.running_mean * bn.weight / (bn.running_var + bn.eps) ** 0.5 |
| 78 | m = nn.ConvTranspose2d( |
| 79 | w.size(1), w.size(0), w.shape[2:], stride=self.c.stride, |
| 80 | padding=self.c.padding, dilation=self.c.dilation, groups=self.c.groups) |
| 81 | m.weight.data.copy_(w) |
| 82 | m.bias.data.copy_(b) |
| 83 | return m |
| 84 | |
| 85 | |
| 86 | def stem_b4_transpose(in_chs, out_chs, activation): |