20240519 Colab 에서는 onnxruntime cuda provider 가 동작하지 않는 거 같다.

import torch
import torchvision
import numpy as np
import onnxruntime

print(f"onnxruntime.get_device(): {onnxruntime.get_device()}")

model = torchvision.models.resnet18(weights=None)
model.eval()

input_on_cpu = torch.randn((1, 3, 224, 224), dtype=torch.float32)
torch.onnx.export(model, input_on_cpu, "resnet18.onnx", input_names=['input'], output_names=['output'])

session = onnxruntime.InferenceSession("resnet18.onnx", providers=['CUDAExecutionProvider'])
print(f"session.get_providers(): {session.get_providers()}")

 

위 코드를 Colab에서 실행시켜봤다.

 

 

GPU로 추론되려면 session.get_providers()를 호출했을때 CUDAExe... 가 나와야하는데 CPUExe... 만 나오고 안나온다.