Spaces:
Runtime error
Runtime error
| # Copyright 2023 The TensorFlow 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. | |
| """Tests for multitask.task_sampler.""" | |
| import tensorflow as tf, tf_keras | |
| from official.modeling.multitask import configs | |
| from official.modeling.multitask import task_sampler as sampler | |
| class TaskSamplerTest(tf.test.TestCase): | |
| def setUp(self): | |
| super(TaskSamplerTest, self).setUp() | |
| self._task_weights = {'A': 1.0, 'B': 2.0, 'C': 3.0} | |
| def test_uniform_sample_distribution(self): | |
| uniform_sampler = sampler.get_task_sampler( | |
| configs.TaskSamplingConfig(type='uniform'), self._task_weights) | |
| for step in range(5): | |
| cumulative_distribution = uniform_sampler.task_cumulative_distribution( | |
| tf.constant(step, dtype=tf.int64)) | |
| self.assertAllClose([0.333333, 0.666666, 1.0], | |
| cumulative_distribution.numpy()) | |
| def test_proportional_sample_distribution(self): | |
| prop_sampler = sampler.get_task_sampler( | |
| configs.TaskSamplingConfig( | |
| type='proportional', | |
| proportional=configs.ProportionalSampleConfig(alpha=2.0)), | |
| self._task_weights) | |
| # CucmulativeOf(Normalize([1.0^2, 2.0^2, 3.0^2])) | |
| for step in range(5): | |
| cumulative_distribution = prop_sampler.task_cumulative_distribution( | |
| tf.constant(step, dtype=tf.int64)) | |
| self.assertAllClose([0.07142857, 0.35714286, 1.0], | |
| cumulative_distribution.numpy()) | |
| def test_annealing_sample_distribution(self): | |
| num_epoch = 3 | |
| step_per_epoch = 6 | |
| annel_sampler = sampler.get_task_sampler( | |
| configs.TaskSamplingConfig( | |
| type='annealing', | |
| annealing=configs.AnnealingSampleConfig( | |
| steps_per_epoch=step_per_epoch, | |
| total_steps=step_per_epoch * num_epoch)), self._task_weights) | |
| global_step = tf.Variable( | |
| 0, dtype=tf.int64, name='global_step', trainable=False) | |
| expected_cumulative_epochs = [[0.12056106, 0.4387236, 1.0], | |
| [0.16666667, 0.5, 1.0], | |
| [0.22477472, 0.5654695, 1.0]] | |
| for epoch in range(num_epoch): | |
| for _ in range(step_per_epoch): | |
| cumulative_distribution = annel_sampler.task_cumulative_distribution( | |
| tf.constant(global_step, dtype=tf.int64)) | |
| global_step.assign_add(1) | |
| self.assertAllClose(expected_cumulative_epochs[epoch], | |
| cumulative_distribution.numpy()) | |
| if __name__ == '__main__': | |
| tf.test.main() | |