def _run(init_dir, dataset, model, loss_fn, optimizer, rank):
with tempfile.TemporaryDirectory() as tmpdir:
distributed.initialize(init_dir, rank, 2)
output = os.path.join(tmpdir, "out")
train_supervised(output,
dataset,
model,
optimizer,
loss_fn,
MockEvaluate("key"),
"key",
collate.collate,
1,
Epoch(2),
n_dataloader_worker=0)
if rank == 0:
assert os.path.exists(os.path.join(output, "snapshot_iter_2"))
assert os.path.exists(os.path.join(output, "log.json"))
with open(os.path.join(output, "log.json")) as file:
log = json.load(file)
assert isinstance(log, list)
assert 1 == len(log)
assert 1 == len(os.listdir(os.path.join(output, "model")))
assert os.path.exists(os.path.join(output, "model.pt"))
assert os.path.exists(os.path.join(output, "optimizer.pt"))
else:
assert not os.path.exists(os.path.join(output, "snapshot_iter_2"))
assert not os.path.exists(os.path.join(output, "log.json"))
assert not os.path.exists(os.path.join(output, "model"))
return model.state_dict(), optimizer.state_dict()
def _run(self, init_dir, input, output, model, synthesizer, dataset, rank):
distributed.initialize(init_dir, rank, 2)
evaluate(
input, output, dataset,
model, synthesizer,
{
"accuracy": use_environment(
Accuracy(),
in_keys=["actual", ["ground_truth", "expected"]],
value_key="actual",
),
"bleu": use_environment(
Bleu(),
in_keys=["actual", ["ground_truth", "expected"]],
value_key="actual",
),
}
)
def launch(config_file: str, option: Optional[str], tmpdir: str,
rank: Optional[int], n_process: Optional[int]):
logging.set_level(L.INFO)
logger.info(f"Launch config file: {config_file}")
configs = load_config(config_file)
if option == "test":
logger.info("Modify configs for testing")
configs = modify_config_for_test(configs, tmpdir)
elif option == "profile":
logger.info("Modify configs for profiling")
output_dir = configs["output_dir"] # TODO
configs = modify_config_for_profile(configs, tmpdir)
distributed.initialize(tmpdir, rank, n_process)
rank = distributed.rank()
seed = random.randint(0, 2**31) + rank
logger.info(f"Fix seed={seed}")
rng = np.random.RandomState(seed)
torch.manual_seed(rng.randint(0, 2**32 - 1))
np.random.seed(rng.randint(0, 2**32 - 1))
random.seed(rng.randint(0, 2**32 - 1))
logger.info("Run main")
if option == "profile":
cprofile = cProfile.Profile()
cprofile.enable()
with profile() as torch_prof:
parse_config(configs)["/main"]
cprofile.disable()
torch.save(torch_prof,
os.path.join(output_dir, f"torch_profiler-{rank}.pt"))
cprofile.dump_stats(os.path.join(output_dir, f"cprofile-{rank}.pt"))
else:
parse_config(configs)["/main"]
def _run(init_dir, rank):
distributed.initialize(init_dir, rank, 2)
return [
x.xs for x in distributed.all_gather(
Values([i for i in range(2 * rank + 1)]))
]
def _run(self, init_dir, dataset, metrics, rank):
distributed.initialize(init_dir, rank, 2)
return EvaluateSynthesizer(dataset, synthesize,
metrics=metrics)()