def __init__(self, num_segments, dyn_a=True, dyn_b=True):
super(UnsortedSegmentMinDynNet, self).__init__()
self.unsorted_segment_min = P.UnsortedSegmentMin()
self.gpu_convert_to_dynamic_shape = inner.GpuConvertToDynamicShape()
self.num_segments = num_segments
self.to_dyn_1 = dyn_a
self.to_dyn_2 = dyn_b
def __init__(self, num_segments):
super(UnsortedSegmentMinNet, self).__init__()
self.unsorted_segment_min = P.UnsortedSegmentMin()
self.num_segments = num_segments
def test_3d_single_init():
context.set_context(mode=context.GRAPH_MODE, device_target='GPU')
input_x = Tensor(np.arange(4 * 5 * 3, dtype=np.float32).reshape(4, 5, 3),
dtype=mindspore.float32)
segment_ids = Tensor([3, 0, 1, -1], mstype.int32)
net = P.UnsortedSegmentMin()
num_segments = 4
output = net(input_x, segment_ids, num_segments).asnumpy()
expect = np.array([[[1.5000000e+01, 1.6000000e+01, 1.7000000e+01],
[1.8000000e+01, 1.9000000e+01, 2.0000000e+01],
[2.1000000e+01, 2.2000000e+01, 2.3000000e+01],
[2.4000000e+01, 2.5000000e+01, 2.6000000e+01],
[2.7000000e+01, 2.8000000e+01, 2.9000000e+01]],
[[3.0000000e+01, 3.1000000e+01, 3.2000000e+01],
[3.3000000e+01, 3.4000000e+01, 3.5000000e+01],
[3.6000000e+01, 3.7000000e+01, 3.8000000e+01],
[3.9000000e+01, 4.0000000e+01, 4.1000000e+01],
[4.2000000e+01, 4.3000000e+01, 4.4000000e+01]],
[[3.4028235e+38, 3.4028235e+38, 3.4028235e+38],
[3.4028235e+38, 3.4028235e+38, 3.4028235e+38],
[3.4028235e+38, 3.4028235e+38, 3.4028235e+38],
[3.4028235e+38, 3.4028235e+38, 3.4028235e+38],
[3.4028235e+38, 3.4028235e+38, 3.4028235e+38]],
[[0.0000000e+00, 1.0000000e+00, 2.0000000e+00],
[3.0000000e+00, 4.0000000e+00, 5.0000000e+00],
[6.0000000e+00, 7.0000000e+00, 8.0000000e+00],
[9.0000000e+00, 1.0000000e+01, 1.1000000e+01],
[1.2000000e+01, 1.3000000e+01,
1.4000000e+01]]]).astype(np.float32)
np.testing.assert_array_almost_equal(output, expect)
num_segments = 6
output = net(input_x, segment_ids, num_segments).asnumpy()
expect = np.array([[[1.5000000e+01, 1.6000000e+01, 1.7000000e+01],
[1.8000000e+01, 1.9000000e+01, 2.0000000e+01],
[2.1000000e+01, 2.2000000e+01, 2.3000000e+01],
[2.4000000e+01, 2.5000000e+01, 2.6000000e+01],
[2.7000000e+01, 2.8000000e+01, 2.9000000e+01]],
[[3.0000000e+01, 3.1000000e+01, 3.2000000e+01],
[3.3000000e+01, 3.4000000e+01, 3.5000000e+01],
[3.6000000e+01, 3.7000000e+01, 3.8000000e+01],
[3.9000000e+01, 4.0000000e+01, 4.1000000e+01],
[4.2000000e+01, 4.3000000e+01, 4.4000000e+01]],
[[3.4028235e+38, 3.4028235e+38, 3.4028235e+38],
[3.4028235e+38, 3.4028235e+38, 3.4028235e+38],
[3.4028235e+38, 3.4028235e+38, 3.4028235e+38],
[3.4028235e+38, 3.4028235e+38, 3.4028235e+38],
[3.4028235e+38, 3.4028235e+38, 3.4028235e+38]],
[[0.0000000e+00, 1.0000000e+00, 2.0000000e+00],
[3.0000000e+00, 4.0000000e+00, 5.0000000e+00],
[6.0000000e+00, 7.0000000e+00, 8.0000000e+00],
[9.0000000e+00, 1.0000000e+01, 1.1000000e+01],
[1.2000000e+01, 1.3000000e+01, 1.4000000e+01]],
[[3.4028235e+38, 3.4028235e+38, 3.4028235e+38],
[3.4028235e+38, 3.4028235e+38, 3.4028235e+38],
[3.4028235e+38, 3.4028235e+38, 3.4028235e+38],
[3.4028235e+38, 3.4028235e+38, 3.4028235e+38],
[3.4028235e+38, 3.4028235e+38, 3.4028235e+38]],
[[3.4028235e+38, 3.4028235e+38, 3.4028235e+38],
[3.4028235e+38, 3.4028235e+38, 3.4028235e+38],
[3.4028235e+38, 3.4028235e+38, 3.4028235e+38],
[3.4028235e+38, 3.4028235e+38, 3.4028235e+38],
[3.4028235e+38, 3.4028235e+38,
3.4028235e+38]]]).astype(np.float32)
np.testing.assert_array_almost_equal(output, expect)
def __init__(self, strategy1, strategy2, num_segments):
super(Net, self).__init__()
self.virtual_dataset = _VirtualDataset()
self.merge_op = P.UnsortedSegmentMin().shard((strategy1, strategy2))
self.num_segments = num_segments