| 6 | |
| 7 | |
| 8 | class MaxPool3dSamePadding(nn.MaxPool3d): |
| 9 | |
| 10 | def compute_pad(self, dim, s): |
| 11 | if s % self.stride[dim] == 0: |
| 12 | return max(self.kernel_size[dim] - self.stride[dim], 0) |
| 13 | else: |
| 14 | return max(self.kernel_size[dim] - (s % self.stride[dim]), 0) |
| 15 | |
| 16 | def forward(self, x): |
| 17 | # compute 'same' padding |
| 18 | (batch, channel, t, h, w) = x.size() |
| 19 | pad_t = self.compute_pad(0, t) |
| 20 | pad_h = self.compute_pad(1, h) |
| 21 | pad_w = self.compute_pad(2, w) |
| 22 | |
| 23 | pad_t_f = pad_t // 2 |
| 24 | pad_t_b = pad_t - pad_t_f |
| 25 | pad_h_f = pad_h // 2 |
| 26 | pad_h_b = pad_h - pad_h_f |
| 27 | pad_w_f = pad_w // 2 |
| 28 | pad_w_b = pad_w - pad_w_f |
| 29 | |
| 30 | pad = (pad_w_f, pad_w_b, pad_h_f, pad_h_b, pad_t_f, pad_t_b) |
| 31 | x = F.pad(x, pad) |
| 32 | return super(MaxPool3dSamePadding, self).forward(x) |
| 33 | |
| 34 | |
| 35 | class Unit3D(nn.Module): |