wehub-resource-sync
122 lines
3.9 KiB
Python
122 lines
3.9 KiB
Python
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import random
|
|
import unittest
|
|
|
|
import numpy as np
|
|
|
|
import paddle
|
|
import paddle.distributed as dist
|
|
from paddle.distributed import fleet
|
|
|
|
|
|
def set_random_seed(seed):
|
|
"""Set random seed for reproducibility."""
|
|
random.seed(seed)
|
|
np.random.seed(seed)
|
|
paddle.seed(seed)
|
|
fleet.meta_parallel.model_parallel_random_seed(seed)
|
|
|
|
|
|
def class_center_sample_numpy(label, classes_list, num_samples):
|
|
unique_label = np.unique(label)
|
|
nranks = len(classes_list)
|
|
class_interval = np.cumsum(np.insert(classes_list, 0, 0))
|
|
pos_class_center_per_device = []
|
|
unique_label_per_device = []
|
|
|
|
for i in range(nranks):
|
|
index = np.logical_and(
|
|
unique_label >= class_interval[i],
|
|
unique_label < class_interval[i + 1],
|
|
)
|
|
pos_class_center_per_device.append(
|
|
unique_label[index] - class_interval[i]
|
|
)
|
|
unique_label_per_device.append(unique_label[index])
|
|
|
|
num_samples_per_device = []
|
|
for pos_class_center in pos_class_center_per_device:
|
|
num_samples_per_device.append(max(len(pos_class_center), num_samples))
|
|
sampled_class_interval = np.cumsum(np.insert(num_samples_per_device, 0, 0))
|
|
|
|
remapped_dict = {}
|
|
for i in range(nranks):
|
|
for idx, v in enumerate(
|
|
unique_label_per_device[i], sampled_class_interval[i]
|
|
):
|
|
remapped_dict[v] = idx
|
|
|
|
remapped_label = []
|
|
for l in label:
|
|
remapped_label.append(remapped_dict[l])
|
|
|
|
return remapped_label, pos_class_center_per_device
|
|
|
|
|
|
class TestParallelClassCenterSampleOp(unittest.TestCase):
|
|
def setUp(self):
|
|
strategy = fleet.DistributedStrategy()
|
|
fleet.init(is_collective=True, strategy=strategy)
|
|
|
|
def test_class_center_sample(self):
|
|
rank_id = dist.get_rank()
|
|
nranks = dist.get_world_size()
|
|
|
|
seed = 1025
|
|
set_random_seed(seed)
|
|
paddle.seed(rank_id * 10)
|
|
random.seed(seed)
|
|
np.random.seed(seed)
|
|
|
|
batch_size = 20
|
|
num_samples = 6
|
|
|
|
for dtype in ('int32', 'int64'):
|
|
for _ in range(5):
|
|
classes_list = np.random.randint(10, 15, (nranks,))
|
|
num_class = np.sum(classes_list)
|
|
|
|
np_label = np.random.randint(
|
|
0, num_class, (batch_size,), dtype=dtype
|
|
)
|
|
label = paddle.to_tensor(np_label, dtype=dtype)
|
|
(
|
|
np_remapped_label,
|
|
np_sampled_class_center_per_device,
|
|
) = class_center_sample_numpy(
|
|
np_label, classes_list, num_samples
|
|
)
|
|
(
|
|
remapped_label,
|
|
sampled_class_index,
|
|
) = paddle.nn.functional.class_center_sample(
|
|
label, classes_list[rank_id], num_samples
|
|
)
|
|
np.testing.assert_allclose(
|
|
remapped_label.numpy(), np_remapped_label
|
|
)
|
|
np_sampled_class_index = np_sampled_class_center_per_device[
|
|
rank_id
|
|
]
|
|
np.testing.assert_allclose(
|
|
sampled_class_index.numpy()[: len(np_sampled_class_index)],
|
|
np_sampled_class_index,
|
|
)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|