Spaces:
Build error
Build error
File size: 1,164 Bytes
fa8453f |
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 |
import onnx
import onnxruntime
device_types_list = ["cpu", "cuda"]
available_providers = onnxruntime.get_available_providers()
def get_device_and_provider(device='cpu'):
options = onnxruntime.SessionOptions()
options.log_severity_level = 3
if device == 'cuda':
if "CUDAExecutionProvider" in available_providers:
provider = [("CUDAExecutionProvider", {"cudnn_conv_algo_search": "DEFAULT"}), "CPUExecutionProvider"]
options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
else:
device = 'cpu'
provider = ["CPUExecutionProvider"]
else:
device = 'cpu'
provider = ["CPUExecutionProvider"]
return device, provider, options
data_type_bytes = {'uint8': 1, 'int8': 1, 'uint16': 2, 'int16': 2, 'float16': 2, 'float32': 4}
def estimate_max_batch_size(resolution, chunk_size=1024, data_type='float32', channels=3):
pixel_size = data_type_bytes.get(data_type, 1)
image_size = resolution[0] * resolution[1] * pixel_size * channels
number_of_batches = (chunk_size * 1024 * 1024) // image_size
return max(number_of_batches, 1)
|