一、PyTorch C++ API 概覽

PyTorch 的 C++ API 提供了一個(gè)強(qiáng)大的工具集，用于在 C++ 環(huán)境中進(jìn)行張量計(jì)算和深度學(xué)習(xí)模型開發(fā)。它主要包括以下幾個(gè)部分：

1.1 ATen 庫(kù)

ATen 是 PyTorch 的基礎(chǔ)張量庫(kù)，提供了豐富的張量操作和數(shù)學(xué)運(yùn)算功能。

#include <ATen/ATen.h>


at::Tensor a = at::ones({2, 2}, at::kInt);
at::Tensor b = at::randn({2, 2});
auto c = a + b.to(at::kInt);

1.2 Autograd 自動(dòng)求導(dǎo)

Autograd 是 PyTorch C++ API 的自動(dòng)微分組件，擴(kuò)展了 ATen 的功能，使其支持自動(dòng)求導(dǎo)。

#include <torch/torch.h>


torch::Tensor a = torch::ones({2, 2}, torch::requires_grad());
torch::Tensor b = torch::randn({2, 2});
auto c = a + b;
c.backward();

1.3 C++ 前端

C++ 前端提供了高層接口，用于構(gòu)建和訓(xùn)練神經(jīng)網(wǎng)絡(luò)模型。

#include <torch/torch.h>


class SimpleModel : public torch::nn::Module {
public:
    SimpleModel() {
        linear = register_module("linear", torch::nn::Linear(10, 2));
    }


    torch::Tensor forward(torch::Tensor x) {
        return linear(x);
    }


private:
    torch::nn::Linear linear;
};


int main() {
    SimpleModel model;
    torch::Tensor input = torch::randn({1, 10});
    torch::Tensor output = model(input);
    return 0;
}

1.4 TorchScript 支持

TorchScript 是 PyTorch 的 JIT 編譯器和解釋器，支持模型的序列化和優(yōu)化。

#include <torch/torch.h>


int main() {
    // 加載 TorchScript 模型
    torch::jit::script::Module model;
    model.load("model.pt");


    // 執(zhí)行模型推理
    torch::Tensor input = torch::randn({1, 3, 224, 224});
    torch::Tensor output = model.forward({input});


    return 0;
}

1.5 C++ 擴(kuò)展

C++ 擴(kuò)展允許開發(fā)者通過自定義 C++ 和 CUDA 代碼擴(kuò)展 PyTorch 的功能。

#include <torch/torch.h>


torch::Tensor custom_add(torch::Tensor x, torch::Tensor y) {
    return x + y;
}


PYBIND11_MODULE(custom_ops, m) {
    m.def("custom_add", &custom_add, "A custom add operation");
}

二、開發(fā)環(huán)境搭建

2.1 安裝 PyTorch C++ API

可以從 PyTorch 官方網(wǎng)站獲取安裝包，或通過源代碼編譯。

## 使用 conda 安裝
conda install pytorch torchvision torchaudio pytorch-cpp -c pytorch

2.2 配置開發(fā)工具

推薦使用支持 C++17 的編譯器，如 GCC 9 或更高版本。同時(shí)，可以使用 CMake 來管理項(xiàng)目構(gòu)建。

cmake_minimum_required(VERSION 3.18)
project(MyPyTorchProject)


find_package(Torch REQUIRED)


add_executable(my_app main.cpp)
target_link_libraries(my_app PRIVATE ${TORCH_LIBRARIES})

三、模型開發(fā)與訓(xùn)練

3.1 定義模型

使用 torch::nn::Module 定義神經(jīng)網(wǎng)絡(luò)模型。

#include <torch/torch.h>


class MyModel : public torch::nn::Module {
public:
    MyModel() {
        conv1 = register_module("conv1", torch::nn::Conv2d(
            torch::nn::Conv2dOptions(1, 20, 5)
        ));
        fc1 = register_module("fc1", torch::nn::Linear(20 * 20 * 20, 500));
        fc2 = register_module("fc2", torch::nn::Linear(500, 10));
    }


    torch::Tensor forward(torch::Tensor x) {
        x = torch::functional::ReLU(conv1(x));
        x = torch::max_pool2d(x, 2);
        x = x.view({-1, 20 * 20 * 20});
        x = torch::functional::ReLU(fc1(x));
        x = fc2(x);
        return x;
    }


private:
    torch::nn::Conv2d conv1;
    torch::nn::Linear fc1, fc2;
};

3.2 數(shù)據(jù)加載與處理

使用 torch::data 加載和處理數(shù)據(jù)。

#include <torch/data.h>
#include <torch/datasets.h>


using namespace torch::data;


class MyDataset : public torch::data::Dataset<MyDataset> {
public:
    MyDataset(std::string path) : path_(std::move(path)) {}


    torch::data::Example<> get(size_t index) override {
        // 實(shí)現(xiàn)數(shù)據(jù)加載邏輯
        torch::Tensor data = ...;
        torch::Tensor target = ...;
        return {data, target};
    }


    torch::optional<size_t> size() const override {
        return 1000; // 數(shù)據(jù)集大小
    }


private:
    std::string path_;
};


int main() {
    auto dataset = MyDataset("data");
    auto dataloader = make_data_loader(
        dataset,
        DataLoaderOptions().batch_size(32).workers(4)
    );


    for (auto& batch : *dataloader) {
        auto data = batch.data;
        auto target = batch.target;
        // 訓(xùn)練邏輯
    }


    return 0;
}

3.3 模型訓(xùn)練與優(yōu)化

使用優(yōu)化器進(jìn)行模型訓(xùn)練。

#include <torch/optim.h>


int main() {
    MyModel model;
    auto optimizer = torch::optim::SGD(
        model.parameters(),
        torch::optim::SGDOptions(0.01).momentum(0.9)
    );


    for (auto& batch : *dataloader) {
        auto data = batch.data;
        auto target = batch.target;


        optimizer.zero_grad();
        auto output = model(data);
        auto loss = torch::nn::functional::nll_loss(output, target);
        loss.backward();
        optimizer.step();
    }


    return 0;
}

四、模型推理與部署

4.1 模型保存與加載

保存和加載模型參數(shù)或整個(gè)模型。

#include <torch/serialize.h>


int main() {
    MyModel model;


    // 保存模型
    torch::save(model, "model.pth");


    // 加載模型
    MyModel loaded_model;
    torch::load(loaded_model, "model.pth");


    return 0;
}

4.2 TorchScript 模型推理

加載并運(yùn)行 TorchScript 模型。

#include <torch/jit.h>


int main() {
    // 加載 TorchScript 模型
    torch::jit::script::Module model;
    model.load("model.pt");


    // 模型推理
    torch::Tensor input = torch::randn({1, 3, 224, 224});
    std::vector<torch::jit::IValue> inputs;
    inputs.push_back(input);
    torch::Tensor output = model.forward(inputs).toTensor();


    return 0;
}

五、性能優(yōu)化技巧

5.1 使用混合精度訓(xùn)練

在訓(xùn)練過程中使用混合精度加速計(jì)算。

#include <torchcuda.h>


int main() {
    MyModel model;
    model.cuda();


    auto scaler = torch::cuda::amp::GradScaler();
    for (auto& batch : *dataloader) {
        auto data = batch.data.cuda();
        auto target = batch.target.cuda();


        scaler.scale(loss).backward();
        scaler.step(optimizer);
        scaler.update();
    }


    return 0;
}

5.2 多 GPU 并行訓(xùn)練

使用多 GPU 進(jìn)行模型并行訓(xùn)練。

#include <torch/distributed.h>
#include <torch/data/distributed.h>


int main() {
    // 初始化分布式環(huán)境
    torch::distributed::init_process_group(torch::distributed::Backend::NCCL, std::string("env://"));


    int rank = torch::distributed::get_rank();
    int world_size = torch::distributed::get_world_size();


    MyModel model;
    model.cuda(rank);


    // 數(shù)據(jù)并行
    auto model_ddp = torch::nn::DataParallel(model);


    return 0;
}

六、總結(jié)與展望

PyTorch C++ API 提供了強(qiáng)大的功能，使開發(fā)者能夠在 C++ 環(huán)境中高效地進(jìn)行深度學(xué)習(xí)模型的開發(fā)和部署。通過合理利用 ATen、Autograd、C++ 前端、TorchScript 和 C++ 擴(kuò)展，可以構(gòu)建高性能的機(jī)器學(xué)習(xí)應(yīng)用。

關(guān)注編程獅（W3Cschool）平臺(tái)，獲取更多 PyTorch C++ API 開發(fā)相關(guān)的教程和案例。