def conv3d(
input,
pointwise,
spatial,
bias=None,
stride=1,
padding=0,
dilation=1,
groups=1,
):
stride = triple(stride)
padding = triple(padding)
dilation = triple(dilation)
_ = F.conv3d(input, pointwise, bias, 1, 0, 1, groups)
for i, weight in enumerate(spatial):
stri = one_diff_tuple(3, 1, stride[i], i)
pad = one_diff_tuple(3, 0, padding[i], i)
dil = one_diff_tuple(3, 1, dilation[i], i)
_ = F.conv3d(_, weight, None, stri, pad, dil, _.shape[1])
return _
def _to_tuple(self, module: nn.Module, val: Union[Tuple[int],
int]) -> Tuple[int]:
if isinstance(val, tuple):
return val
module_name = module.__class__.__name__.lower()
if "1d" in module_name:
return single(val)
elif "2d" in module_name:
return double(val)
elif "3d" in module_name:
return triple(val)
else:
raise ValueError(
f"Couldn't infer tuple size for class {module.__class__.__name__}. "
"Please pass an explicit tuple.")
def __init__(
self,
in_channels,
out_channels,
kernel_size,
factorized=False,
bn_depth=None,
bn_spatial=None,
repeats=1,
bn_repeats=1,
stride=2,
padding=0,
groups=1,
bias=False,
padding_mode="zeros",
checkpoint=False,
):
kernel_size = triple(kernel_size)
same_pad = tuple([x // 2 for x in kernel_size])
if bn_depth is not None or bn_spatial is not None:
if factorized:
repeated = BottleneckFactorized3d(
in_channels,
in_channels,
kernel_size,
bn_depth,
bn_spatial,
1,
same_pad,
1,
bn_repeats,
groups,
bias,
padding_mode,
checkpoint,
)
else:
repeated = Bottleneck3d(
in_channels,
in_channels,
kernel_size,
bn_depth,
bn_spatial,
1,
same_pad,
1,
bn_repeats,
groups,
bias,
padding_mode,
checkpoint,
)
else:
if factorized:
repeated = Conv3d(
in_channels,
in_channels,
kernel_size,
1,
same_pad,
1,
groups,
bias,
padding_mode,
)
else:
repeated = nn.Conv3d(
in_channels,
in_channels,
kernel_size,
1,
same_pad,
1,
groups,
bias,
padding_mode,
)
if factorized:
final = Conv3d(
in_channels,
out_channels,
kernel_size,
stride,
padding,
1,
groups,
bias,
padding_mode,
)
else:
final = nn.Conv3d(
in_channels,
out_channels,
kernel_size,
stride,
padding,
1,
groups,
bias,
padding_mode,
)
super(DownSample3d, self).__init__(repeated, final, repeats)