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)))