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