def test_avg_pool_nchw(self): pb = PB({'attr': { 'data_format': PB({ 's': b"NCHW" }), 'strides': PB({ 'list': PB({ "i": self.strides }) }), 'ksize': PB({ 'list': PB({ "i": self.ksize }) }), 'padding': PB({ 's': b'VALID' }) }}) self.expected = { 'window': np.array(self.ksize, dtype=np.int8), 'spatial_dims': [2, 3], 'stride': np.array(self.strides, dtype=np.int8), 'pool_method': "avg", } node = PB({'pb': pb}) AvgPoolFrontExtractor.extract(node) self.res = node self.res["infer"](None) self.call_args = self.infer_mock.call_args self.expected_call_args = (None, "avg") self.compare()
def test_pool_defaults(self): pb = PB({'attr': { 'data_format': PB({ 's': b"NHWC" }), 'strides': PB({ 'list': PB({ "i": self.strides }) }), 'ksize': PB({ 'list': PB({"i": self.ksize}) }), 'padding': PB({ 's': b'VALID' }) }}) self.expected = { 'pad': None, # will be inferred when input shape is known 'pad_spatial_shape': None, 'type': 'Pooling', 'exclude_pad': 'true', } node = PB({'pb': pb}) AvgPoolFrontExtractor.extract(node) self.res = node self.res["infer"](None) self.call_args = self.infer_mock.call_args self.expected_call_args = (None, None) self.compare()