NVIDIA-AI-IOT / trt_pose

Real-time pose estimation accelerated with NVIDIA TensorRT
MIT License
974 stars 291 forks source link

if only CPU, can run trt_pose project ? #134

Open lian-yeh opened 3 years ago

lian-yeh commented 3 years ago

hello, my computer have no GPU, I run live_demo.py error message :
raise AssertionError("Torch not compiled with CUDA enabled") AssertionError: Torch not compiled with CUDA enabled

how to solve it ? thank you

lweicker commented 2 years ago

Yes it is possible to run this project only with CPU. Assuming you have installed trt_pose and use densenet121_baseline_att_256x256_B model, here follows a snippet to load the model with cpu.

import json

import torch
import trt_pose.models

HUMAN_POSE_JSON_PATH = 'human_pose.json'
MODEL_PATH = 'my_model.pth'

def _load_topology(human_pose_path: str):
    with open(human_pose_path, 'r') as f:
        human_pose = json.load(f)
    skeleton = human_pose['skeleton']
    keypoints = human_pose['keypoints']
    length_skeleton = len(skeleton)
    topology = torch.zeros((length_skeleton, 4)).int()
    for k in range(length_skeleton):
        topology[k][0] = 2 * k
        topology[k][1] = 2 * k + 1
        topology[k][2] = skeleton[k][0] - 1
        topology[k][3] = skeleton[k][1] - 1
    return topology, length_skeleton, len(keypoints)

def _load_model(model_path: str, num_links: int, num_parts: int):
    model = trt_pose.models.densenet121_baseline_att(num_parts, 2 * num_links,
                                                     pretrained=False)
    model = model.eval()
    model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
    return model

topology, num_links, num_parts = _load_topology(HUMAN_POSE_JSON_PATH)
model = _load_model(MODEL_PATH, num_links, num_parts)

Regarding preprocess, you'll have to replace, in live_demo.ipynb, this line: device = torch.device('cuda') by device = torch.device('cpu').