Unified model loader from CLI args.
(args)
| 159 | |
| 160 | |
| 161 | def load_model(args) -> nn.Module: |
| 162 | """Unified model loader from CLI args.""" |
| 163 | dtype = _parse_dtype(args.dtype) |
| 164 | |
| 165 | if args.model: |
| 166 | print(f"Loading model from file: {args.model} (class: {args.class_name})") |
| 167 | model = load_model_from_file(args.model, args.class_name) |
| 168 | elif args.module: |
| 169 | print(f"Loading model from module: {args.module} (class: {args.class_name})") |
| 170 | extra_kwargs = {} |
| 171 | if dtype == torch.float16: |
| 172 | extra_kwargs["torch_dtype"] = torch.float16 |
| 173 | elif dtype == torch.bfloat16: |
| 174 | extra_kwargs["torch_dtype"] = torch.bfloat16 |
| 175 | model = load_model_from_module( |
| 176 | args.module, args.class_name, pretrained=args.pretrained, **extra_kwargs |
| 177 | ) |
| 178 | else: |
| 179 | raise ValueError("Must specify either --model (file path) or --module (Python module)") |
| 180 | |
| 181 | model = model.to(dtype=dtype) |
| 182 | |
| 183 | if torch.cuda.is_available(): |
| 184 | try: |
| 185 | model = model.cuda() |
| 186 | except RuntimeError as e: |
| 187 | if "out of memory" in str(e).lower(): |
| 188 | print(f"WARNING: OOM moving model to GPU. Trying with smaller footprint...") |
| 189 | torch.cuda.empty_cache() |
| 190 | model = model.half().cuda() |
| 191 | else: |
| 192 | raise |
| 193 | |
| 194 | model.eval() |
| 195 | return model |
| 196 | |
| 197 | |
| 198 | # --------------------------------------------------------------------------- |
no test coverage detected