| 6 | |
| 7 | class DynamicConv2d(nn.Module): |
| 8 | def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, groups=1): |
| 9 | super().__init__() |
| 10 | self.in_channels = in_channels |
| 11 | self.out_channels = out_channels |
| 12 | self.kernel_size = kernel_size |
| 13 | self.stride = stride |
| 14 | self.padding = padding |
| 15 | self.groups = groups |
| 16 | |
| 17 | self.weight = nn.Parameter(torch.randn(out_channels, in_channels // groups, kernel_size, kernel_size)) |
| 18 | self.bias = nn.Parameter(torch.zeros(out_channels)) |
| 19 | |
| 20 | nn.init.kaiming_normal_(self.weight, mode='fan_out', nonlinearity='relu') |
| 21 | |
| 22 | def forward(self, x): |
| 23 | return F.conv2d(x, self.weight, self.bias, self.stride, self.padding, groups=self.groups) |