Spaces:
Sleeping
Sleeping
| import datetime | |
| import argparse, importlib | |
| from pytorch_lightning import seed_everything | |
| import torch | |
| import torch.distributed as dist | |
| def setup_dist(local_rank): | |
| if dist.is_initialized(): | |
| return | |
| torch.cuda.set_device(local_rank) | |
| torch.distributed.init_process_group('nccl', init_method='env://') | |
| def get_dist_info(): | |
| if dist.is_available(): | |
| initialized = dist.is_initialized() | |
| else: | |
| initialized = False | |
| if initialized: | |
| rank = dist.get_rank() | |
| world_size = dist.get_world_size() | |
| else: | |
| rank = 0 | |
| world_size = 1 | |
| return rank, world_size | |
| if __name__ == '__main__': | |
| now = datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S") | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--module", type=str, help="module name", default="inference") | |
| parser.add_argument("--local_rank", type=int, nargs="?", help="for ddp", default=0) | |
| args, unknown = parser.parse_known_args() | |
| inference_api = importlib.import_module(args.module, package=None) | |
| inference_parser = inference_api.get_parser() | |
| inference_args, unknown = inference_parser.parse_known_args() | |
| seed_everything(inference_args.seed) | |
| setup_dist(args.local_rank) | |
| torch.backends.cudnn.benchmark = True | |
| rank, gpu_num = get_dist_info() | |
| print("@CoLVDM Inference [rank%d]: %s"%(rank, now)) | |
| inference_api.run_inference(inference_args, gpu_num, rank) |