torch::nn::Module is the C++ LibTorch base class for defining neural-network models. A custom network inherits from it, implements layers, and exposes forward computation for training or inference. section 17.5.1
A simple feed-forward network
A Net class can model MNIST digits with flattened 28×28 inputs, a 128-neuron hidden layer using ReLU, and 10 output neurons using log-softmax. Training calls model->forward, computes loss, backpropagates, and updates model->parameters(). section 17.5.1
// Python
model = torch.nn.Linear(10, 5)
input = torch.randn(1, 10)
output = model(input)
// Equivalent C++
auto model = torch::nn::Linear(10, 5);
auto input = torch::randn({1, 10});
auto output = model->forward(input);
Compare equivalent Python and C++ APIs for constructing a linear module, creating input, and calling forward. section 17.5.1
Listing 17.12 Training in C++ with LibTorch
void model_train_example() {
auto model = std::make_shared<Net>(); #1
torch::optim::SGD optimizer(model->parameters(), /*lr=*/0.01);
// Create fake data
auto x = torch::randn({64, 784});
auto y = torch::randint(0, 10, {64});
model->train();
for (size_t epoch = 0; epoch < 3; ++epoch) { #2
auto prediction = model->forward(x);
auto loss = torch::nll_loss(prediction, y);
optimizer.zero_grad();
loss.backward();
optimizer.step();
std::cout << "Epoch: " << epoch << " | Loss: " << loss.item<float>()
<< std::endl;
}
// Inference
model->eval();
torch::NoGradGuard no_grad; #3
auto test_input = torch::randn({1, 784});
auto output = model->forward(test_input);
auto predicted = output.argmax(1);
std::cout << "Prediction: " << predicted.item<int64_t>() << std::endl;
}
Follow a complete C++ training and inference loop, including optimizer steps, evaluation mode, and gradient disabling. section 17.5.1
