第 0 章 环境准备与第一次运行

torch::Device select_device() {
    if (torch::cuda::is_available()) return torch::Device(torch::kCUDA);
    return torch::Device(torch::kCPU);
}

torch::Device device = select_device();            // 优先 CUDA，无 GPU 时回落 CPU
torch::manual_seed(123);                           // 固定随机种子
at::globalContext().setAllowTF32CuBLAS(false);     // 教学约束：禁用 TF32

auto x = torch::tensor({{1.0f, 2.0f, 3.0f},
                        {4.0f, 5.0f, 6.0f}});       // (2,3) Float
auto w = torch::tensor({{1.0f}, {1.0f}, {1.0f}});   // (3,1)

auto y_mat  = x.matmul(w);                          // (2,3)@(3,1) = (2,1)
auto y_soft = torch::softmax(x, /*dim=*/1);         // 沿最后一维归一化
auto y_mask = x.masked_fill(x > 3.0f, -1e9);        // 条件填充 → 注意力掩码的雏形
auto y_topk = std::get<0>(x.topk(/*k=*/2, /*dim=*/1));  // 每行 top-2 → 解码策略用

auto row = x.index({1});                             // 第 2 行（下标从 0 起）→ 形状 (3)
auto col = x.index({torch::indexing::Slice(), 2});   // 每行第 3 列 → 形状 (2)

auto xg = torch::tensor({1.0, 2.0, 3.0}, torch::kFloat).requires_grad_(true);
auto yg = (xg * xg).sum();          // y = Σx²，dy/dx = 2x
yg.backward();
std::cout << xg.grad();             // [2, 4, 6]

class DemoModuleImpl : public torch::nn::Module {
 public:
    DemoModuleImpl(int64_t in_dim, int64_t hidden_dim) {
        fc_      = register_module("fc", torch::nn::Linear(in_dim, hidden_dim));
        dropout_ = register_module("drop", torch::nn::Dropout(0.1));
        scale_   = register_parameter("scale", torch::tensor(1.0).reshape({1}));
        bias_    = register_buffer("bias", torch::arange(0, hidden_dim, torch::kFloat));
    }
    torch::Tensor forward(torch::Tensor x) {
        auto h = torch::relu(fc_->forward(x));
        h = h * scale_;                       // 广播：(B,H) * (1,)
        h = h.masked_fill(h <= 0, -1e9);
        return h.mean(/*dim=*/-1, /*keepdim=*/true);
    }
 private:
    torch::nn::Linear fc_{nullptr};
    torch::nn::Dropout dropout_{nullptr};
    torch::Tensor scale_;
    torch::Tensor bias_;
};
TORCH_MODULE(DemoModule);   // 生成 ModuleHolder<DemoModuleImpl>

for (const auto& p : model->parameters()) { ... }          // 优化器入口
for (const auto& np : model->named_parameters(/*recurse=*/true)) { ... }  // 名字 + 形状
for (const auto& nm : model->named_modules()) { ... }      // 子模块树

auto opt = torch::optim::AdamW(
    model->parameters(),
    torch::optim::AdamWOptions(1e-3).betas({0.9, 0.95}).weight_decay(0.01));

opt.zero_grad();                       // 1. 清零梯度
auto loss = ((logits - target).pow(2)).mean();   // 2. 前向 + 算 loss
loss.backward();                       // 3. 反向传播
opt.step();                            // 4. 更新参数
// 5.（循环回到 1）

opt.param_groups().front().options().set_lr(1e-4);

torch::save(model, "/tmp/demo.pt");
torch::load(model_loaded, "/tmp/demo.pt");