diff --git a/README.md b/README.md index 3114bab..4c2d3a8 100644 --- a/README.md +++ b/README.md @@ -12,12 +12,35 @@ $ which construe $ construe --help ``` +There are several top-level configurations that you can specify either as an environment variable or a command line option before the command. The environment variables are as follows: + +- `$CONSTRUE_ENV` or `$ENV`: specify the name of the experimental environment for comparison purposes. +- `$CONSTRUE_DEVICE` or `$TORCH_DEVICE`: specify the name of the default device to use with PyTorch e.g. cpu, mps, or cuda. + +The command line utility help is as follows: + +``` +Usage: construe [OPTIONS] COMMAND [ARGS]... + +Options: + --version Show the version and exit. + -d, --device TEXT specify the pytorch device to run on e.g. cpu, mps or + cuda + -e, --env TEXT name of the experimental environment for comparison + (default is hostname) + -h, --help Show this message and exit. + +Commands: + basic + moondream +``` + ## Basic Benchmarks The basic benchmarks implement dot product benchmarks from the [PyTorch documentation](https://pytorch.org/tutorials/recipes/recipes/benchmark.html). These benchmarks can be run using `construe basic`; for example by running: ``` -$ construe basic -e "MacBook Pro 2022 M1" -o results-macbook.pickle +$ construe -e "MacBook Pro 2022 M1" basic -o results-macbook.pickle ``` The `-e` flag specifies the environment for comparison purposes and the `-o` flag saves the measurements out to disk as a Pickle file that can be loaded for comparison to other environments later. diff --git a/construe/__main__.py b/construe/__main__.py index 3ff7f75..e0334e3 100644 --- a/construe/__main__.py +++ b/construe/__main__.py @@ -3,7 +3,10 @@ """ import click +import torch +import platform +from click import ClickException from .version import get_version from .basic import BasicBenchmark from .moondream import MoonDreamBenchmark @@ -16,17 +19,40 @@ @click.group(context_settings=CONTEXT_SETTINGS) @click.version_option(get_version(), message="%(prog)s v%(version)s") -def main(): - pass - - -@main.command() +@click.option( + "-d", + "--device", + default=None, + type=str, + envvar=["CONSTRUE_DEVICE", "TORCH_DEVICE"], + help="specify the pytorch device to run on e.g. cpu, mps or cuda", +) @click.option( "-e", "--env", default=None, + envvar=["CONSTRUE_ENV", "ENV"], help="name of the experimental environment for comparison (default is hostname)", ) +@click.pass_context +def main(ctx, env=None, device=None): + if device is not None: + try: + torch.set_default_device(device) + except RuntimeError as e: + raise ClickException(str(e)) + + click.echo(f"using torch.device(\"{device}\")") + + if env is None: + env = platform.node() + + ctx.ensure_object(dict) + ctx.obj["device"] = device + ctx.obj["env"] = env + + +@main.command() @click.option( "-o", "--saveto", @@ -53,18 +79,23 @@ def main(): type=int, help="set the random seed for random generation", ) -def basic(**kwargs): +@click.pass_context +def basic(ctx, **kwargs): + kwargs["env"] = ctx.obj["env"] benchmark = BasicBenchmark(**kwargs) benchmark.run() @main.command() -def moondream(**kwargs): +@click.pass_context +def moondream(ctx, **kwargs): + kwargs["env"] = ctx["env"] benchmark = MoonDreamBenchmark(**kwargs) benchmark.run() if __name__ == "__main__": main( + obj={}, prog_name="construe", ) diff --git a/construe/basic.py b/construe/basic.py index cd3fb0e..055e204 100644 --- a/construe/basic.py +++ b/construe/basic.py @@ -4,9 +4,9 @@ See: https://pytorch.org/tutorials/recipes/recipes/benchmark.html """ +import tqdm import torch import pickle -import platform import torch.utils.benchmark as benchmark from itertools import product @@ -35,9 +35,6 @@ def __init__(self, env=None, saveto=None, num_threads=None, fuzz=False, seed=Non if num_threads is None: num_threads = torch.get_num_threads() - if env is None: - env = platform.node() - self.env = env self.saveto = saveto self.num_threads = num_threads @@ -46,7 +43,7 @@ def __init__(self, env=None, saveto=None, num_threads=None, fuzz=False, seed=Non def run(self): results = [] - dataset = self.fuzzer().take(10) if self.fuzz else self.static() + dataset = self.fuzzer().take(10) if self.fuzz else list(self.static()) kwargs = { "label": "Batched Dot", @@ -54,7 +51,7 @@ def run(self): "env": self.env, } - for tensors, tensor_params, params in dataset: + for tensors, tensor_params, params in tqdm.tqdm(dataset, leave=False): sub_label = f"{params['k0']:<6} x {params['k1']:<4} {'' if tensor_params['x']['is_contiguous'] else '(discontiguous)'}" # noqa results.append( benchmark.Timer( diff --git a/construe/moondream.py b/construe/moondream.py index ec29c16..28fff79 100644 --- a/construe/moondream.py +++ b/construe/moondream.py @@ -27,7 +27,7 @@ def __init__(self): @profile def run(self): results = [] - for path in tqdm.tqdm(self.dataset): + for path in tqdm.tqdm(self.dataset, leave=False): encoded, encode_time = self.moondream.encode_image(path) inference_time = self.moondream.inference(encoded) results.append((encode_time, inference_time))