| from __future__ import annotations |
| from pathlib import Path |
| import numpy as np |
| import torch |
| import onnxruntime as ort |
|
|
| _NP_DTYPES={ |
| 'tensor(float)':np.float32,'tensor(float16)':np.float16,'tensor(double)':np.float64, |
| 'tensor(int64)':np.int64,'tensor(int32)':np.int32,'tensor(bool)':np.bool_, |
| } |
| _TORCH_DTYPES={np.float32:torch.float32,np.float16:torch.float16,np.float64:torch.float64,np.int64:torch.int64,np.int32:torch.int32,np.bool_:torch.bool} |
|
|
| class TorchORTCudaSession: |
| """Static-shape ONNX Runtime CUDA wrapper with zero-copy PyTorch I/O binding.""" |
| def __init__(self,path:str|Path): |
| so=ort.SessionOptions(); so.graph_optimization_level=ort.GraphOptimizationLevel.ORT_ENABLE_ALL |
| providers=['CUDAExecutionProvider','CPUExecutionProvider'] |
| self.session=ort.InferenceSession(str(path),sess_options=so,providers=providers) |
| if self.session.get_providers()[0] != 'CUDAExecutionProvider': |
| raise RuntimeError(f'CUDAExecutionProvider unavailable: {self.session.get_providers()}') |
| self.inputs={x.name:x for x in self.session.get_inputs()}; self.outputs={x.name:x for x in self.session.get_outputs()} |
| def run(self,feeds:dict[str,torch.Tensor]): |
| io=self.session.io_binding(); keep=[] |
| for name,meta in self.inputs.items(): |
| t=feeds[name].contiguous(); npdt=_NP_DTYPES[meta.type]; tdt=_TORCH_DTYPES[npdt] |
| if t.dtype!=tdt: t=t.to(dtype=tdt) |
| if not t.is_cuda: t=t.cuda(non_blocking=True) |
| feeds[name]=t; keep.append(t) |
| io.bind_input(name,'cuda',0,npdt,tuple(t.shape),int(t.data_ptr())) |
| outs={} |
| for name,meta in self.outputs.items(): |
| shape=tuple(int(x) for x in meta.shape); npdt=_NP_DTYPES[meta.type]; y=torch.empty(shape,device='cuda',dtype=_TORCH_DTYPES[npdt]); keep.append(y); outs[name]=y |
| io.bind_output(name,'cuda',0,npdt,shape,int(y.data_ptr())) |
| self.session.run_with_iobinding(io); io.synchronize_outputs(); return outs |
|
|
| class DinoORT(torch.nn.Module): |
| def __init__(self,path): super().__init__(); self.engine=TorchORTCudaSession(path); self.device=torch.device('cuda') |
| def forward(self,pixel_values,is_training=True): return {'x_prenorm':self.engine.run({'pixel_values':pixel_values})['x_prenorm']} |
| def to(self,*a,**k):return self |
| def cpu(self):return self |
| def eval(self):return self |
|
|
| class DsineORT(torch.nn.Module): |
| def __init__(self,path): super().__init__(); self.engine=TorchORTCudaSession(path); self.device=torch.device('cuda') |
| def forward(self,image,intrins): return [self.engine.run({'image':image,'intrins':intrins})['normal']] |
| def to(self,*a,**k):return self |
| def cpu(self):return self |
| def eval(self):return self |
|
|
