def test_pool_module_meta_kwargs():
kernel_size = (2, )
x = make_pool_module(nn.MaxPool1d, kernel_size=kernel_size)
actual = meta.pool_module_meta(x, kernel_size=kernel_size)["kernel_size"]
desired = kernel_size
assert actual == desired
def test_main(self):
pool_module_meta = {
"kernel_size": (2, ),
"stride": (3, ),
"padding": (4, ),
}
x = make_pool_module(nn.MaxPool1d, **pool_module_meta)
actual = meta.pool_module_meta(x)
desired = pool_module_meta
assert actual == desired
def test_pool_module_meta(self):
pool_module_meta = {
"kernel_size": (2, ),
"stride": (3, ),
"padding": (4, ),
}
x = nn.MaxPool1d(**pool_module_meta)
actual = meta.pool_module_meta(x)
desired = pool_module_meta
self.assertDictEqual(actual, desired)
def multi_layer_encoder(impl_params: bool = True, ) -> enc.MultiLayerEncoder:
r"""Multi-layer encoder from :cite:`GEB2016`.
Args:
impl_params: If ``True``, the :class:`~torch.nn.MaxPool2d` in
the ``multi_layer_encoder`` are exchanged for :class:`~torch.nn.AvgPool2d`.
"""
# TODO: check if the encoder used inplace ops
multi_layer_encoder_ = enc.vgg19_multi_layer_encoder(
framework="caffe", internal_preprocessing=False, allow_inplace=True)
if impl_params:
return multi_layer_encoder_
for name, module in multi_layer_encoder_.named_children():
if isinstance(module, nn.MaxPool2d):
multi_layer_encoder_._modules[name] = nn.AvgPool2d(
**meta.pool_module_meta(module))
return multi_layer_encoder_
def _pool_guide(module: PoolModule, guide: torch.Tensor) -> torch.Tensor:
return cast(torch.Tensor, F.max_pool2d(guide, **pool_module_meta(module)))
