Jetson NX 配置 pytorch的问题及解决方法

下面我将介绍如何在Jetson NX上配置PyTorch，并提供两个示例说明。

Jetson NX配置PyTorch的问题

由于Jetson NX使用的是ARM架构，而PyTorch官方只提供了x86和AMD64架构下的二进制包，所以我们需要手动编译安装PyTorch，或使用第三方提供的二进制包来进行安装。另外，需要注意的是，Jetson NX上需要使用具有ARM64架构的PyTorch版本才能够正常工作。

解决方法1：手动编译安装PyTorch

1. 先安装依赖

sudo apt-get update
sudo apt-get install libopenblas-dev libblas-dev m4 cmake cython python3-dev python3-yaml python3-setuptools

1. 下载PyTorch源码，并进入源码目录

git clone --recursive https://github.com/pytorch/pytorch
cd pytorch

1. 执行脚本开始编译并安装PyTorch，过程可能需要较长时间

export NO_CUDA=1
export NO_DISTRIBUTED=1
export NO_MKLDNN=1
python3 setup.py build
sudo python3 setup.py install

1. 编译完成后，可以输入以下命令测试PyTorch是否安装成功

python3 -c 'import torch; print(torch.__version__)'

解决方法2：使用第三方提供的二进制包

1. 从NVIDIA提供的JetPack下载页面（https://developer.nvidia.com/jetpack-sdk）下载与你的Jetson NX版本匹配的JetPack安装包。

2. 使用以下命令安装JetPack：

sudo apt-mark hold openssh-server
sudo apt-get update
sudo apt-get install curl
curl https://oaci4teeo.qnssl.com/Jetpack-L4T-4.6.1-linux-aarch64.run -k -o Jetpack-L4T-4.6.1-linux-aarch64.run
sudo chmod +x Jetpack-L4T-4.6.1-linux-aarch64.run
sudo ./Jetpack-L4T-4.6.1-linux-aarch64.run

过程中需要输入NVIDIA账号进行验证。

1. 安装完JetPack后，在JetPack的组件列表中找到并安装PyTorch。

2. 安装完成后，可以输入以下命令测试PyTorch是否安装成功：

python3 -c 'import torch; print(torch.__version__)'

示例1：使用PyTorch进行图片分类

1. 安装torchvision库

sudo pip3 install torchvision

1. 下载并解压imagenet数据集，可以使用以下命令

wget https://storage.googleapis.com/download.tensorflow.org/models/tflite_11_05_08/tf_2.4.0/imagenet_saved_model.tar.gz
tar xvzf imagenet_saved_model.tar.gz

1. 创建一个Python脚本image_classification.py，使用以下代码进行图片分类：

import torch
import torchvision
import torchvision.transforms as transforms

# 加载图片数据集
valset = torchvision.datasets.ImageFolder(
    root='./imagenet_saved_model/Imagenet/', transform=transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                             std=[0.229, 0.224, 0.225])
    ]))

# 加载训练好的模型
model = torchvision.models.resnet50(pretrained=True)

# 开启评估模式
model.eval()

# 预测每张图片的类别
with torch.no_grad():
    for i, (inputs, targets) in enumerate(valset):
        outputs = model(inputs.unsqueeze(0))
        _, predicted = torch.max(outputs.data, 1)
        print('Predicted: ', valset.classes[predicted.item()])

1. 使用以下命令执行image_classification.py

python3 image_classification.py

示例2：使用PyTorch进行手写数字识别

1. 安装MNIST数据集

sudo pip3 install python-mnist

1. 创建一个Python脚本mnist_classification.py，使用以下代码进行手写数字识别：

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.datasets as datasets
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from torchsummary import summary
from mnist import MNIST

# 加载MNIST数据集
mndata = MNIST('./mnist_data')
train_images, train_labels = mndata.load_training()
test_images, test_labels = mndata.load_testing()

# 转换成PyTorch所需的格式
train_images = torch.Tensor(train_images)
train_labels = torch.LongTensor(train_labels)
test_images = torch.Tensor(test_images)
test_labels = torch.LongTensor(test_labels)

train_set = list(zip(train_images, train_labels))
test_set = list(zip(test_images, test_labels))

# 定义模型结构
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(784, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = x.view(-1, 784)
        x = nn.functional.relu(self.fc1(x))
        x = self.fc2(x)
        return nn.functional.log_softmax(x, dim=1)

# 定义超参数
batch_size = 64
learning_rate = 0.001
epochs = 10

# 加载数据集
train_loader = DataLoader(train_set, batch_size=batch_size)
test_loader = DataLoader(test_set, batch_size=batch_size)

# 定义模型和优化器
model = Net()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

# 输出模型结构
summary(model, (1, 28, 28))

# 训练模型
for epoch in range(epochs):
    model.train()
    for i, (inputs, targets) in enumerate(train_loader):
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = nn.functional.nll_loss(outputs, targets)
        loss.backward()
        optimizer.step()

    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for i, (inputs, targets) in enumerate(test_loader):
            outputs = model(inputs)
            _, predicted = torch.max(outputs.data, 1)
            total += targets.size(0)
            correct += (predicted == targets).sum().item()

    accuracy = 100.0 * correct / total
    print('Epoch [{}/{}], Loss: {:.4f}, Accuracy: {:.2f}%'.format(epoch+1, epochs, loss.item(), accuracy))

1. 使用以下命令执行mnist_classification.py

python3 mnist_classification.py

希望这个攻略对你有所帮助！