Sample tasks according to task weights as well as training progress. See http://proceedings.mlr.press/v97/stickland19a/stickland19a.pdf
| 82 | |
| 83 | |
| 84 | class AnnealingTaskSampler(TaskSampler): |
| 85 | """Sample tasks according to task weights as well as training progress. |
| 86 | |
| 87 | See http://proceedings.mlr.press/v97/stickland19a/stickland19a.pdf |
| 88 | """ |
| 89 | |
| 90 | def __init__(self, |
| 91 | task_weights: Dict[Text, Union[float, int]], |
| 92 | steps_per_epoch: int, |
| 93 | total_steps: int): |
| 94 | super(AnnealingTaskSampler, self).__init__(task_weights=task_weights) |
| 95 | self._steps_per_epoch = tf.cast(steps_per_epoch, dtype=tf.float32) |
| 96 | self._total_epochs = tf.cast( |
| 97 | total_steps / self._steps_per_epoch, dtype=tf.float32) |
| 98 | |
| 99 | def task_cumulative_distribution(self, global_step: tf.Tensor) -> tf.Tensor: |
| 100 | cur_epoch = tf.math.floor( |
| 101 | tf.cast(global_step, dtype=tf.float32) / self._steps_per_epoch) |
| 102 | alpha = 1.0 - 0.8 * (cur_epoch - 1) / (self._total_epochs - 1 + 1e-10) |
| 103 | task_weight_dict_ordered_list = [ |
| 104 | weight for _, weight in self._task_weights.items() |
| 105 | ] |
| 106 | task_sizes = tf.math.pow( |
| 107 | tf.constant(task_weight_dict_ordered_list, dtype=tf.float32), |
| 108 | tf.cast(alpha, dtype=tf.float32)) |
| 109 | dynamic_task_distribution = task_sizes / tf.reduce_sum(task_sizes) |
| 110 | return tf.math.cumsum(dynamic_task_distribution) |
| 111 | |
| 112 | |
| 113 | def get_task_sampler(config: configs.TaskSamplingConfig, |