Skip to content

Commit

Permalink
Fix util paths
Browse files Browse the repository at this point in the history
  • Loading branch information
szaman19 committed Feb 16, 2024
1 parent 60bbdfd commit a4a853b
Show file tree
Hide file tree
Showing 4 changed files with 64 additions and 186 deletions.
105 changes: 60 additions & 45 deletions applications/FLASK/Transformer/arg_utils.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
import argparse
import data_utils
import dataset_utils


# Command-line arguments
Expand All @@ -8,70 +8,85 @@ def add_transformer_architecture_arguments(args: argparse.Namespace):
Adds the command line arguments to specify transformer architecture model
parameters. This is only relevant for the encoder-decoder transformer model.
"""
args.add_argument('--num-attention-heads',
action='store',
default=8,
type=int,
help='number of parallel attention layers (default: 8)',
metavar='NUM')
args.add_argument('--embed-dim',
action='store',
default=512,
type=int,
help='embedding space dimension (default: 512)',
metavar='NUM')
args.add_argument('--feedforward-dim',
action='store',
default=0,
type=int,
help='feedforward network dimension. If zero, set to be '
'4 times the embedding dimension (default: 0)',
metavar='NUM')
args.add_argument('--num-layers',
action='store',
default=6,
type=int,
help='Number of encoder and decoder layers (default: 6)',
metavar='NUM')
args.add_argument(
"--num-attention-heads",
action="store",
default=8,
type=int,
help="number of parallel attention layers (default: 8)",
metavar="NUM",
)
args.add_argument(
"--embed-dim",
action="store",
default=512,
type=int,
help="embedding space dimension (default: 512)",
metavar="NUM",
)
args.add_argument(
"--feedforward-dim",
action="store",
default=0,
type=int,
help="feedforward network dimension. If zero, set to be "
"4 times the embedding dimension (default: 0)",
metavar="NUM",
)
args.add_argument(
"--num-layers",
action="store",
default=6,
type=int,
help="Number of encoder and decoder layers (default: 6)",
metavar="NUM",
)


def add_dataset_arguments(args: argparse.Namespace, default: str):
"""
Adds dataset-related arguments to an existing argparse object.
"""
args.add_argument('--dataset',
type=str,
default=default,
help=f'Which dataset to use (default: {default})',
choices=data_utils.available_datasets())
args.add_argument('--dataset-fraction',
action='store',
default=1.0,
type=float,
help='Fraction of dataset to use (default: 1.0)',
metavar='NUM')
args.add_argument(
"--dataset",
type=str,
default=default,
help=f"Which dataset to use (default: {default})",
choices=dataset_utils.available_datasets(),
)
args.add_argument(
"--dataset-fraction",
action="store",
default=1.0,
type=float,
help="Fraction of dataset to use (default: 1.0)",
metavar="NUM",
)


def add_training_arguments(parser: argparse.ArgumentParser):
parser.add_argument("--skip-validation",
action="store_true",
default=False,
help="Do not run validation (default: false)")
parser.add_argument(
"--skip-validation",
action="store_true",
default=False,
help="Do not run validation (default: false)",
)
parser.add_argument(
"--always-shuffle",
action="store_true",
default=False,
help=
"Always shuffle training dataset, even if pretokenized (default: false)"
help="Always shuffle training dataset, even if pretokenized (default: false)",
)
parser.add_argument(
"--validation-set-fraction",
type=float,
default=0.01,
help="Fraction of the validation dataset to use (default: 0.001)")
help="Fraction of the validation dataset to use (default: 0.001)",
)
parser.add_argument(
"--save-prototext",
action="store_true",
default=False,
help="Save prototext experiment file instead of protobin (slower but "
"debuggable) (default: false)")
"debuggable) (default: false)",
)
21 changes: 0 additions & 21 deletions applications/FLASK/Transformer/dataset_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,24 +24,3 @@ def load_dataset(name: str):
"""
sys.path.append(dataset_dir)
return importlib.import_module(name)


def add_dataset_arguments(args: argparse.Namespace, default: str):
"""
Adds dataset-related arguments to an existing argparse object.
"""
args.add_argument(
"--dataset",
type=str,
default=default,
help=f"Which dataset to use (default: {default})",
choices=available_datasets(),
)
args.add_argument(
"--dataset-fraction",
action="store",
default=1.0,
type=float,
help="Fraction of dataset to use (default: 1.0)",
metavar="NUM",
)
47 changes: 0 additions & 47 deletions applications/FLASK/Transformer/network.py
Original file line number Diff line number Diff line change
Expand Up @@ -177,50 +177,3 @@ def _add_encoder_decoder_loss(
# Compute cross-entropy
ce = lbann.CrossEntropy(preds, labels, use_labels=True)
return lbann.Scale(ce, constant=1 / (sequence_length - 1))


def add_transformer_architecture_arguments(args: argparse.Namespace):
"""
Adds the command line arguments to specify transformer architecture model
parameters. This is only relevant for the encoder-decoder transformer model.
"""
args.add_argument(
"--num-attention-heads",
action="store",
default=8,
type=int,
help="number of parallel attention layers (default: 8)",
metavar="NUM",
)
args.add_argument(
"--embed-dim",
action="store",
default=512,
type=int,
help="embedding space dimension (default: 512)",
metavar="NUM",
)
args.add_argument(
"--feedforward-dim",
action="store",
default=0,
type=int,
help="feedforward network dimension. If zero, set to be "
"4 times the embedding dimension (default: 0)",
metavar="NUM",
)
args.add_argument(
"--num-layers",
action="store",
default=6,
type=int,
help="Number of encoder and decoder layers (default: 6)",
metavar="NUM",
)
args.add_argument(
"--positional-encoding",
type=str,
default="learned",
help="The type of positional encoding to use " "(default: learned)",
choices=[s.name.lower() for s in InputEncoding],
)
77 changes: 4 additions & 73 deletions applications/FLASK/Transformer/trainer.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
import paths

import dataset_utils
import arg_utils
import applications.FLASK.Transformer.network as network


Expand Down Expand Up @@ -267,85 +268,15 @@ def make_batch_script(
return script


def add_training_arguments(parser: argparse.ArgumentParser):
parser.add_argument(
"--skip-validation",
action="store_true",
default=False,
help="Do not run validation (default: false)",
)
parser.add_argument(
"--always-shuffle",
action="store_true",
default=False,
help="Always shuffle training dataset, even if pretokenized (default: false)",
)
parser.add_argument(
"--validation-set-fraction",
type=float,
default=0.01,
help="Fraction of the validation dataset to use (default: 0.001)",
)
parser.add_argument(
"--save-prototext",
action="store_true",
default=False,
help="Save prototext experiment file instead of protobin (slower but "
"debuggable) (default: false)",
)
parser.add_argument(
"--validate-every",
type=int,
default=100,
help="Run validation every N steps (default: 100)",
)
parser.add_argument(
"--random-seed",
action="store",
default=None,
type=int,
help=f"Set random seed explicitly",
metavar="NUM",
)
parser.add_argument(
"--pbar-newline-interval",
type=int,
default=100,
help="Number of iterations in progress bar before "
"printing a newline (default: 100)",
)
parser.add_argument(
"--pbar-width",
type=int,
default=30,
help="Progress bar width, if enabled (default: 30)",
)
parser.add_argument(
"--pbar-moving-avg",
type=int,
default=10,
help="Progress bar iteration time moving average "
"length in iterations. Disable moving average with 1 "
"(default: 10)",
)
parser.add_argument(
"--pbar-scientific",
action="store_true",
default=False,
help="Use scientific notation for objective value "
"printouts in progress bar (default: false)",
)


def main():
# Setup command line options
parser = argparse.ArgumentParser()
lbann.contrib.args.add_scheduler_arguments(parser, "lbann_transformer")
lbann.contrib.args.add_profiling_arguments(parser)
lbann.contrib.args.add_training_arguments(parser)
network.add_transformer_architecture_arguments(parser)
add_training_arguments(parser)
dataset_utils.add_dataset_arguments(parser, default="qm9")
arg_utils.add_transformer_architecture_arguments(parser)
arg_utils.add_training_arguments(parser)
arg_utils.add_dataset_arguments(parser, default="qm9")

parser.add_argument(
"--optimizer",
Expand Down

0 comments on commit a4a853b

Please sign in to comment.