60 lines
1.6 KiB
C++
60 lines
1.6 KiB
C++
/**
|
|
* @file pytorch_integration.cpp
|
|
* @brief Example: train a small neural network with FCES via libtorch.
|
|
*/
|
|
|
|
#include "fces/optimizer.hpp"
|
|
#include <iostream>
|
|
#include <torch/torch.h>
|
|
|
|
struct TinyNet : torch::nn::Module {
|
|
torch::nn::Linear fc1{nullptr}, fc2{nullptr};
|
|
|
|
TinyNet() {
|
|
fc1 = register_module("fc1", torch::nn::Linear(10, 32));
|
|
fc2 = register_module("fc2", torch::nn::Linear(32, 1));
|
|
}
|
|
|
|
torch::Tensor forward(torch::Tensor x) {
|
|
x = torch::relu(fc1->forward(x));
|
|
return fc2->forward(x);
|
|
}
|
|
};
|
|
|
|
int main() {
|
|
auto model = std::make_shared<TinyNet>();
|
|
|
|
std::vector<torch::Tensor> params;
|
|
for (auto &p : model->parameters())
|
|
params.push_back(p);
|
|
|
|
fces::FCESOptimizer optimizer(params, fces::FCESConfig{}
|
|
.set_lr(1.6e-3f)
|
|
.set_population_size(200)
|
|
.set_total_steps(1000));
|
|
|
|
// Generate synthetic regression data
|
|
auto x_train = torch::randn({100, 10});
|
|
auto y_train = torch::sin(x_train.sum(1, true));
|
|
|
|
for (int epoch = 0; epoch < 100; ++epoch) {
|
|
optimizer.zero_grad();
|
|
auto pred = model->forward(x_train);
|
|
auto loss = torch::mse_loss(pred, y_train);
|
|
loss.backward();
|
|
optimizer.step();
|
|
optimizer.update_fitness(loss.item<float>());
|
|
|
|
if (epoch % 10 == 0) {
|
|
std::cout << "Epoch " << epoch << " | Loss: " << loss.item<float>()
|
|
<< std::endl;
|
|
}
|
|
}
|
|
|
|
std::cout << "\nTraining complete. Final loss: "
|
|
<< torch::mse_loss(model->forward(x_train), y_train).item<float>()
|
|
<< std::endl;
|
|
|
|
return 0;
|
|
}
|