| """Execution-only CUDA adapter. Frozen routing, weights and FP32 math stay fixed. |
| CPU image transforms/postprocessing/order and native ORT banks remain unchanged. |
| """ |
| import types |
| import torch |
|
|
| def configure_fp32(): |
| assert torch.cuda.is_available(),'CUDA required for the GPU evaluation lane' |
| torch.backends.cuda.matmul.allow_tf32=False;torch.backends.cudnn.allow_tf32=False;torch.backends.cudnn.benchmark=False;torch.set_float32_matmul_precision('highest');torch.backends.cuda.enable_flash_sdp(False);torch.backends.cuda.enable_mem_efficient_sdp(False);torch.backends.cuda.enable_math_sdp(True) |
|
|
| def cpu_preprocess_cuda_forward(self,images,threshold=.4,return_features=False): |
| from surya.common.rfdetr.predictor import IMAGENET_MEAN,IMAGENET_STD |
| import torchvision.transforms.functional as TF |
| if not images:return [] |
| tensors=[];sizes=[] |
| for image in images: |
| image=image.convert('RGB');sizes.append((image.height,image.width));t=TF.to_tensor(image);t=TF.normalize(t,IMAGENET_MEAN,IMAGENET_STD);t=TF.resize(t,(self.resolution,self.resolution));tensors.append(t) |
| with torch.inference_mode(): |
| outputs=self.model(torch.stack(tensors).to('cuda'));features=self._feat.get('f')if return_features else None |
| def cpu(x): |
| if torch.is_tensor(x):return x.cpu() |
| if isinstance(x,dict):return{k:cpu(v)for k,v in x.items()} |
| if isinstance(x,list):return[cpu(v)for v in x] |
| return x |
| outputs=cpu(outputs);features=features.cpu()if features is not None else None;results=self.postprocess(outputs,target_sizes=torch.tensor(sizes));out=[] |
| for i,res in enumerate(results): |
| keep=res['scores']>threshold;record=dict(boxes=res['boxes'][keep],scores=res['scores'][keep],labels=res['labels'][keep]) |
| if return_features and features is not None:record['features']=features[i] |
| out.append(record) |
| return out |
|
|
| def accelerate(engine,reader): |
| configure_fp32();reader.model.to('cuda').eval();reader.device=torch.device('cuda') |
| assert all(p.dtype==torch.float32 for p in reader.model.parameters()) |
| engine.execution_devices=dict(layout_model='cpu:unchanged',RTL_student='cuda:FP32',layout_image_transforms='cpu:unchanged',layout_postprocess='cpu:unchanged',learned_reading_order='cpu:unchanged',native_ONNX_banks='CPUExecutionProvider',TF32=False) |
| return engine |
|
|
