I just use Jetson nano to run resnet34 model to classify fish， but it seems that gpu doesn’t work.
The information is as follows：
torch.cuda.is_available() is TRUE
The code is as follows, please tell me how can I use gpu to accelerate detecting：
import torch
from model import resnet34
from PIL import Image
from torchvision import transforms
import matplotlib.pyplot as plt
import json
import cv2
import time
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("device is %s",device)
data_transform = transforms.Compose(
[transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]) # 预处理
# read class_indict
try:
json_file = open('./class_indices.json', 'r')
class_indict = json.load(json_file)
except Exception as e:
print(e)
exit(-1)
# create model
model = resnet34(num_classes=30)
# load model weights
model_weight_path = "./resNet34.pth"
model.load_state_dict(torch.load(model_weight_path, map_location=device))
model.eval()
cap = cv2.VideoCapture(1)
while cap.isOpened():
success, frame = cap.read()
print("success=" + str(success))
if success:
# 检测
img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) # BGR-->RGB
img = Image.fromarray(img) # 这里ndarray_image为原来的numpy数组类型的输入
# load image
# # [N, C, H, W]
img = data_transform(img)
# expand batch dimension
img = torch.unsqueeze(img, dim=0)
print("start detect")
with torch.no_grad(): # 不对损失梯度进行跟踪
# predict class
output = torch.squeeze(model(img)) # 压缩batch维度
predict = torch.softmax(output, dim=0) # 得到概率分布
predict_cla = torch.argmax(predict).numpy() # argmax寻找最大值对应的索引
print(class_indict[str(predict_cla)], predict[predict_cla].numpy())
text = str(class_indict[str(predict_cla)]) + str(predict[predict_cla].numpy())
print(text)
# 在框的左上角画出标签和置信度
if predict[predict_cla].numpy() > 0.75:
cv2.putText(frame,text,(10, 20),cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 1)
cv2.imshow("capture",frame)
if (cv2.waitKey(1) & 0xFF) == ord('q'):
break
time.sleep(0.1)
cap.release()
cv2.destroyAllWindows()
The same code and project, put on a PC will be normal with gpu inference.
I have no idea.- -!
Can anyone give me some advice? Thank you.
Hi,
How do you install PyTorch on Jetson?
Do you use our prebuilt package?
https://docs.nvidia.com/deeplearning/frameworks/install-pytorch-jetson-platform/index.html
Thanks.
I use this command to install pytorch :
(my_torch) nvidia@nano:~$ pip3 install torch-1.8.0-cp36-cp36m-linux_aarch64.whl
(my_torch) nvidia@nano:~$ python3
Python 3.6.15 | packaged by conda-forge | (default, Dec 3 2021, 19:12:04)
[GCC 9.4.0] on linux
Type “help”, “copyright”, “credits” or “license” for more information.
import torch
print(torch.version)
1.8.0
torch.cuda.is_available()
True
And in the same virtual environment, running yolov5 accelerates normally.
Hello, I use the prebuilt package Python 3.6 - [`torch-1.8.0-cp36-cp36m-linux_aarch64.whl and try another version Python 3.6 - torch-1.6.0-cp36-cp36m-linux_aarch64.whl, it still doesn’t work.
I upload the program, can anyone help me to run this code in Jetson nano?
Maybe just change this video device number and run python predict.py.
cap = cv2.VideoCapture(1)
Thank you.
resnet_tensorboard_fish_Project_Jetson.zip (75.6 MB)
Hi,
For JetPack 4.6.x, the PyTorch version should be 1.11.0a0+17540c5+nv22.01.
Our prebuilt usually contains the nv tag. Please give it a try.
https://developer.download.nvidia.cn/compute/redist/jp/v461/pytorch/
Thanks.
Thank you, how can you find the pytorch package?
I thought that Jetson nano doesn’t support version over 1.10 because of python version should be more than 3.8.
And I install the torch1.11, but I can’t import it.
It makes some error. How can I fix it?
I have fixed the problem by sudo apt-get install libomp-dev, but meet another.
Pytorch is installed, and then I install torchvision0.12.0 as followed picture.
However, when I import torchvison, it occurs an error.
I find that the reason is python version, Nano with 3.6 but torchvision0.12.0 requires over 3.7.
Can you give me some advice? Thank you very much.
Hello, I try it again and find this part log.
And I think Pillow had been installed, it’s no use.
Thanks for your reply.
OK, so PyTorch 1.10 was the last version of PyTorch to officially support Python 3.6. And torchvision v0.11 was the version of torchvision for PyTorch 1.10. So I would recommend uninstalling torchvision and re-installing torchvision v0.11.1 and see if that helps.
Hello, I re-install torchvision0.11.1 successfully, but it’s no use to accelerate inference.
And I am going to try jetson inference project, change resnet model to onnx file. Can you give me some reference? Cause I don’t know how to use onnx_export.py, and use onnx file to predict.
Thx.
@814928072 for Jetson Nano, you should downgrade to PyTorch 1.10, as that was the last PyTorch version to officially support Python 3.6. You can find the wheel in this topic: PyTorch for Jetson
If you follow the PyTorch part of the Hello AI World tutorial from jetson-inference, it covers how to do the training, export the model to ONNX, and then run it with TensorRT: https://github.com/dusty-nv/jetson-inference/blob/master/docs/pytorch-transfer-learning.md
Thank you for your advice.
I think downgrading PyTorch 1.10 with torchvision0.11.1 still doesn’t work cause I have tried so many version such as 1.8,1.6 and in PC PyTorch 1.6 with torchvision 0.7.0 it works well. I find python version in PC is 3.7.9 and guess the key to the problem is python version.
Now I am trying Jetson inference and got a little clue.
OK yes, if you keep having problems with installing PyTorch/torchvision yourself, I recommend trying the l4t-pytorch or jetson-inference containers which already have these components pre-installed for you.
This topic was automatically closed 14 days after the last reply. New replies are no longer allowed.