| | import torch.nn as nn |
| | import torch.nn.functional as F |
| |
|
| | orderings = [ |
| | [0, 1, 3, 4, 5], |
| | [1, 2, 0, 4, 5], |
| | [2, 3, 1, 4, 5], |
| | [3, 0, 2, 4, 5], |
| | [4, 1, 3, 2, 0], |
| | [5, 1, 3, 0, 2], |
| | ] |
| | rotations = [ |
| | [0, 0, 0, 0, 0], |
| | [0, 0, 0,-1, 1], |
| | [0, 0, 0, 2, 2], |
| | [0, 0, 0, 1,-1], |
| | [0, 1,-1, 2, 0], |
| | [0,-1, 1, 0, 2] |
| | ] |
| |
|
| | def _take_right(face, rot): |
| | if rot == 0: |
| | return face[:, :, 0] |
| | elif rot == 1: |
| | return face[:, 0, :].flip(1) |
| | elif rot == 2: |
| | return face[:, :, -1].flip(1) |
| | elif rot == -1: |
| | return face[:, -1, :] |
| |
|
| | def _take_left(face, rot): |
| | if rot == 0: |
| | return face[:, :, -1] |
| | elif rot == 1: |
| | return face[:, -1, :].flip(1) |
| | elif rot == 2: |
| | return face[:, :, 0].flip(1) |
| | elif rot == -1: |
| | return face[:, 0, :] |
| |
|
| | def _take_top(face, rot): |
| | if rot == 0: |
| | return face[:, -1, :] |
| | elif rot == 1: |
| | return face[:, :, 0] |
| | elif rot == 2: |
| | return face[:, 0, :].flip(1) |
| | elif rot == -1: |
| | return face[:, :, -1].flip(1) |
| |
|
| | def _take_bottom(face, rot): |
| | if rot == 0: |
| | return face[:, 0, :] |
| | elif rot == 1: |
| | return face[:, :, -1] |
| | elif rot == 2: |
| | return face[:, -1, :].flip(1) |
| | elif rot == -1: |
| | return face[:, :, 0].flip(1) |
| |
|
| | def valid_pad_conv_fn(x, one_side_pad=False): |
| | if one_side_pad: |
| | x = x[:, :, :-1, :-1] |
| | assert x.ndim == 4 and x.shape[0] == 6 |
| | _, C, H, W = x.shape |
| | y = x.new_empty(6, C, H+2, W+2) |
| | y[..., 1:-1, 1:-1] = x |
| |
|
| | for i in range(6): |
| | r_idx, l_idx, t_idx, b_idx = orderings[i][1:5] |
| | r_rot, l_rot, t_rot, b_rot = rotations[i][1:5] |
| |
|
| | r_edge = _take_right (x[r_idx], r_rot) |
| | l_edge = _take_left (x[l_idx], l_rot) |
| | t_edge = _take_top (x[t_idx], t_rot) |
| | b_edge = _take_bottom(x[b_idx], b_rot) |
| |
|
| | y[i, :, 1:-1, 0 ] = l_edge |
| | y[i, :, 1:-1, -1 ] = r_edge |
| | y[i, :, 0, 1:-1] = t_edge |
| | y[i, :, -1, 1:-1] = b_edge |
| |
|
| | y[i, :, 0, 0 ] = 0.5*(y[i, :, 0, 1] + y[i, :, 1, 0]) |
| | y[i, :, 0, -1 ] = 0.5*(y[i, :, 0, -2] + y[i, :, 1, -1]) |
| | y[i, :, -1, 0 ] = 0.5*(y[i, :, -2, 0] + y[i, :, -1, 1]) |
| | y[i, :, -1,-1 ] = 0.5*(y[i, :, -2, -1] + y[i, :, -1, -2]) |
| |
|
| | if one_side_pad: |
| | return y[:, :, 1:, 1:] |
| |
|
| | return y |
| |
|
| |
|
| | class PaddedConv2d(nn.Conv2d): |
| | def __init__(self, *args, pad_fn=None, one_side_pad=False, **kwargs): |
| | kwargs = dict(kwargs) |
| | kwargs["padding"] = 0 |
| | super().__init__(*args, **kwargs) |
| | self.pad_fn = pad_fn |
| | self.one_side_pad = one_side_pad |
| |
|
| | def forward(self, x): |
| | x = self.pad_fn(x, one_side_pad=self.one_side_pad) |
| | return F.conv2d( |
| | x, self.weight, self.bias, |
| | stride=self.stride, padding=0, |
| | dilation=self.dilation, groups=self.groups |
| | ) |
| |
|
| | @classmethod |
| | def from_existing(cls, conv: nn.Conv2d, pad_fn, one_side_pad=False): |
| | new = cls( |
| | conv.in_channels, conv.out_channels, conv.kernel_size, |
| | stride=conv.stride, padding=0, dilation=conv.dilation, |
| | groups=conv.groups, bias=(conv.bias is not None), |
| | padding_mode="zeros", pad_fn=pad_fn, one_side_pad=one_side_pad |
| | ) |
| | new.weight = conv.weight |
| | if conv.bias is not None: |
| | new.bias = conv.bias |
| | return new |
| |
|
| |
|
| |
|
