File size: 2,542 Bytes
e6aed17 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 | #!/usr/bin/env python3
# Copyright 2026 Xiaomi Corp. (authors: Han Zhu)
#
# See ../../LICENSE for clarification regarding multiple authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Shared utility functions."""
import argparse
import random
import numpy as np
import torch
def str2bool(v):
"""Used in argparse.ArgumentParser.add_argument to indicate
that a type is a bool type and user can enter
- yes, true, t, y, 1, to represent True
- no, false, f, n, 0, to represent False
See https://stackoverflow.com/questions/15008758/parsing-boolean-values-with-argparse # noqa
"""
if isinstance(v, bool):
return v
if v.lower() in ("yes", "true", "t", "y", "1"):
return True
elif v.lower() in ("no", "false", "f", "n", "0"):
return False
else:
raise argparse.ArgumentTypeError("Boolean value expected.")
def get_best_device():
"""Auto-detect the best available device: CUDA > XPU > MPS > CPU."""
if torch.cuda.is_available():
return "cuda"
if hasattr(torch, "xpu") and torch.xpu.is_available():
return "xpu"
if torch.backends.mps.is_available():
return "mps"
return "cpu"
def get_best_device_with_count():
"""Auto-detect best device and return (device_type, device_count)."""
if torch.cuda.is_available():
return "cuda", torch.cuda.device_count()
if hasattr(torch, "xpu") and torch.xpu.is_available():
return "xpu", torch.xpu.device_count()
if torch.backends.mps.is_available():
return "mps", 1
return "cpu", 1
def fix_random_seed(random_seed: int):
"""
Set the same random seed for the libraries and modules.
Includes the ``random`` module, numpy, and torch.
"""
random.seed(random_seed)
np.random.seed(random_seed)
torch.random.manual_seed(random_seed)
# Ensure deterministic ID creation
rd = random.Random()
rd.seed(random_seed)
|