class JsonRunConfiguration(zconf.RunConfig):
# === Required parameters === #
func = zconf.attr(type=str, required=True)
path = zconf.attr(type=str, required=True)
output_base_path = zconf.attr(type=str, required=True)
class Config(zconf.RunConfig):
attr1 = zconf.attr(default=None)
attr2 = zconf.attr(required=True)
class RunConfiguration(zconf.RunConfig):
# === Required parameters === #
task_name = zconf.attr(type=str, required=True)
task_data_dir = zconf.attr(type=str, required=True)
task_config_path = zconf.attr(type=str, required=True)
class RunConfiguration(zconf.RunConfig):
model_config_path = zconf.attr(type=str)
model_tokenizer_path = zconf.attr(type=str)
task_config_base_path = zconf.attr(type=str)
task_cache_base_path = zconf.attr(type=str)
class SingleTaskConfigurator(zconf.RunConfig):
"""Single-task Configurator
Required:
task_name
train_batch_size
(Task config) Need one of:
task_config_path
task_config_base_path
(Task cache) Need one of:
task_cache_path
task_cache_base_path
(Eval batch size) Need one of:
eval_batch_multiplier
eval_batch_size
(Computing max steps) Need one of:
epochs
max_steps
(Set to 0 if not training)
(Task name list) Specify at least one of:
do_train
do_val
do_test
Optional:
gradient_accumulation_steps
eval_subset_num
num_gpus
train_examples_cap
warmup_steps_proportion
"""
task_name = zconf.attr(type=str, default=None)
task_config_base_path = zconf.attr(type=str, default=None)
task_config_path = zconf.attr(type=str, default=None)
task_cache_base_path = zconf.attr(type=str, default=None)
task_cache_path = zconf.attr(type=str, default=None)
do_train = zconf.attr(action="store_true")
do_val = zconf.attr(action="store_true")
do_test = zconf.attr(action="store_true")
train_batch_size = zconf.attr(type=int, required=True)
eval_batch_multiplier = zconf.attr(type=int, default=None)
eval_batch_size = zconf.attr(type=int, default=None)
gradient_accumulation_steps = zconf.attr(type=int, default=1)
eval_subset_num = zconf.attr(type=int, default=500)
train_sample_weights = zconf.attr(type=str, default=None)
train_loss_weights = zconf.attr(type=str, default=None)
weighted_sampling_start_step = zconf.attr(type=int,
default=0) # zero start
weighted_sampling_start_epoch = zconf.attr(type=int,
default=0) # zero start
weighted_loss_start_step = zconf.attr(type=int, default=0) # zero start
weighted_loss_start_epoch = zconf.attr(type=int, default=0) # zero start
no_rewarmup_in_second_stage = zconf.attr(action="store_true")
fix_seed_for_weighted_sampler = zconf.attr(action="store_true")
epochs = zconf.attr(type=int, default=None)
max_steps = zconf.attr(type=int, default=None)
num_gpus = zconf.attr(type=int, default=None)
warmup_steps_proportion = zconf.attr(type=float, default=0.1)
rewarmup_steps_proportion = zconf.attr(type=float, default=0.1)
def create_config(self):
# === Get task config === #
if self.task_config_path:
assert self.task_config_base_path is None
task_config_path = self.task_config_path
elif self.task_config_base_path is not None:
assert self.task_config_path is None
task_config_path = os.path.join(
self.task_config_base_path,
f"{self.task_name}_config.json",
)
else:
raise RuntimeError(
"Require either `task_config_path` or `task_config_base_path`")
# === Get cache === #
if self.task_cache_path is not None:
assert self.task_cache_base_path is None
task_cache_path = self.task_cache_path
elif self.task_cache_base_path is not None:
assert self.task_cache_path is None
task_cache_path = os.path.join(self.task_cache_base_path,
self.task_name)
else:
raise RuntimeError(
"Need `task_cache_path` or `task_cache_base_path`")
task_cache_config = {}
if self.do_train:
task_cache_config["train"] = os.path.join(task_cache_path, "train")
task_cache_config["train_labels"] = os.path.join(
task_cache_path, "train_labels")
if self.do_val:
task_cache_config["val"] = os.path.join(task_cache_path, "val")
task_cache_config["val_labels"] = os.path.join(
task_cache_path, "val_labels")
if self.do_test:
task_cache_config["test"] = os.path.join(task_cache_path, "test")
task_cache_config["test_labels"] = os.path.join(
task_cache_path, "test_labels")
for v in task_cache_config.values():
assert os.path.exists(v)
# === Compute training steps === #
if self.num_gpus:
# We multiply by num_gpus because 1 step is done across (potentially) multiple GPUs
effective_batch_size = (self.train_batch_size *
self.gradient_accumulation_steps *
self.num_gpus)
else:
effective_batch_size = self.train_batch_size * self.gradient_accumulation_steps
logger.info('effective_batch_size: %f', effective_batch_size)
num_examples = get_num_examples_from_cache(
cache_path=os.path.expandvars(task_cache_config["train"]), )
steps_per_epoch = math.ceil(num_examples / effective_batch_size)
if not self.do_train:
assert self.epochs is None
assert self.max_steps is None
max_steps = 0
elif self.max_steps is not None:
max_steps = self.max_steps
elif self.epochs is not None:
assert self.max_steps is None
max_steps = self.epochs * steps_per_epoch
else:
raise RuntimeError("Require either `epochs` or `max_steps`")
if self.weighted_sampling_start_epoch != 0:
if self.weighted_sampling_start_step != 0:
raise ValueError(
'Only one of "weighted-sampling-start-epoch" or "weighted-sampling-start-step" can be specified.'
)
self.weighted_sampling_start_step = self.weighted_sampling_start_epoch * steps_per_epoch
if self.weighted_loss_start_epoch != 0:
if self.weighted_loss_start_step != 0:
raise ValueError(
'Only one of "weighted-loss-start-epoch" or "weighted-loss-start-step" can be specified.'
)
self.weighted_loss_start_step = self.weighted_loss_start_epoch * steps_per_epoch
if self.train_sample_weights is None:
self.weighted_sampling_start_step = 0
if self.train_loss_weights is None:
self.weighted_loss_start_step = 0
# === Compute eval_batch_size === #
if self.eval_batch_size is not None:
assert self.eval_batch_multiplier is None
eval_batch_size = self.eval_batch_size
elif self.eval_batch_multiplier is not None:
assert self.eval_batch_size is None
eval_batch_size = self.train_batch_size * self.eval_batch_multiplier
else:
raise RuntimeError(
"Require either `eval_batch_size` or `eval_batch_multiplier`")
if self.weighted_sampling_start_step != 0 and self.weighted_loss_start_epoch != 0:
raise NotImplementedError()
else:
second_stage_start_steps = self.weighted_sampling_start_step or self.weighted_loss_start_step
first_stage_steps = second_stage_start_steps
warmup_steps = first_stage_steps * self.warmup_steps_proportion
second_stage_steps = max_steps - first_stage_steps
if self.no_rewarmup_in_second_stage:
rewarmup_steps = 0
else:
rewarmup_steps = second_stage_steps * self.rewarmup_steps_proportion
# === Build configuration === #
# Finally, we build our big config dictionary. Congrats!
config_dict = {
"task_config_path_dict": {
self.task_name: task_config_path
},
"task_cache_config_dict": {
self.task_name: task_cache_config
},
"sampler_config": {
"sampler_type": "UniformMultiTaskSampler"
},
"global_train_config": {
"max_steps":
int(max_steps),
"first_stage_steps":
first_stage_steps,
"second_stage_steps":
second_stage_steps,
"warmup_steps":
warmup_steps,
"rewarmup_steps":
rewarmup_steps,
"weighted_sampling_start_step":
self.weighted_sampling_start_step,
"weighted_loss_start_step":
self.weighted_loss_start_step,
"fix_seed_for_weighted_sampler":
self.fix_seed_for_weighted_sampler,
},
"task_specific_configs_dict": {
self.task_name: {
"train_batch_size": self.train_batch_size,
"eval_batch_size": eval_batch_size,
"gradient_accumulation_steps":
self.gradient_accumulation_steps,
"eval_subset_num": self.eval_subset_num,
"train_sample_weights": self.train_sample_weights,
"train_loss_weights": self.train_loss_weights,
}
},
"taskmodels_config": {
"task_to_taskmodel_map": {
self.task_name: self.task_name
},
"taskmodel_config_map": {
self.task_name: None
},
},
"task_run_config": {
"train_task_list": [self.task_name] if self.do_train else [],
"train_val_task_list":
[self.task_name] if self.do_train else [],
"val_task_list": [self.task_name] if self.do_val else [],
"test_task_list": [self.task_name] if self.do_test else [],
},
"metric_aggregator_config": {
"metric_aggregator_type": "EqualMetricAggregator"
},
}
return config_dict
class RunConfiguration(zconf.RunConfig):
# === Required parameters === #
# Directories
task_config_path = zconf.attr(type=str, required=True)
task_cache_base_path = zconf.attr(type=str, required=True)
run_config_path = zconf.attr(type=str, required=True)
output_path = zconf.attr(type=str, required=True)
jiant_path = zconf.attr(type=str, required=True)
model_config = zconf.attr(type=str, required=True)
exp_command_path = zconf.attr(type=str, required=True)
# Others
epochs = zconf.attr(type=int, required=True)
n_trials = zconf.attr(type=int, required=True)
train = zconf.attr(default='cnli',
type=str,
choices={
'cnli', 'cnli_seed', 'snlisub0', 'snlisub1',
'snlisub2', 'snlisub3', 'snlisub4'
},
required=True)
val = zconf.attr(default='mnli',
type=str,
choices={'glue_diagnostic', 'mnli', 'stress'},
required=True)
# === Optional parameters === #
batch_clustering = zconf.attr(action='store_true')
extract_exp_name_valpreds = zconf.attr(action="store_true")
fp16 = zconf.attr(action="store_true")
no_improvements_for_n_evals = zconf.attr(type=int, default=0)
eval_every_steps = zconf.attr(type=int, default=0)
boolq = zconf.attr(action="store_true")
sbatch_name = zconf.attr(default="", type=str)
