| 85 | Default value is 2 (spatial). |
| 86 | """ |
| 87 | def __init__(self, channels, kernel_size, dim=2): |
| 88 | super(AverageSmoothing, self).__init__() |
| 89 | |
| 90 | # Make sure sum of values in gaussian kernel equals 1. |
| 91 | kernel = torch.ones(size=(kernel_size, kernel_size)) / (kernel_size * kernel_size) |
| 92 | |
| 93 | # Reshape to depthwise convolutional weight |
| 94 | kernel = kernel.view(1, 1, *kernel.size()) |
| 95 | kernel = kernel.repeat(channels, *[1] * (kernel.dim() - 1)) |
| 96 | |
| 97 | self.register_buffer('weight', kernel) |
| 98 | self.groups = channels |
| 99 | |
| 100 | if dim == 1: |
| 101 | self.conv = F.conv1d |
| 102 | elif dim == 2: |
| 103 | self.conv = F.conv2d |
| 104 | elif dim == 3: |
| 105 | self.conv = F.conv3d |
| 106 | else: |
| 107 | raise RuntimeError( |
| 108 | 'Only 1, 2 and 3 dimensions are supported. Received {}.'.format(dim) |
| 109 | ) |
| 110 | |
| 111 | def forward(self, input): |
| 112 | """ |