46 lines
1.1 KiB
C++
46 lines
1.1 KiB
C++
#include "fces/optimizer.hpp"
|
|
#include <gtest/gtest.h>
|
|
#include <torch/torch.h>
|
|
|
|
using namespace fces;
|
|
|
|
TEST(OptimizerTest, Construction) {
|
|
auto model = torch::nn::Linear(10, 5);
|
|
std::vector<torch::Tensor> params;
|
|
for (auto &p : model->parameters())
|
|
params.push_back(p);
|
|
|
|
FCESOptimizer opt(params, FCESConfig{}.set_lr(1e-3f));
|
|
EXPECT_EQ(opt.step_count(), 0);
|
|
}
|
|
|
|
TEST(OptimizerTest, StepUpdatesCounter) {
|
|
auto model = torch::nn::Linear(10, 5);
|
|
std::vector<torch::Tensor> params;
|
|
for (auto &p : model->parameters())
|
|
params.push_back(p);
|
|
|
|
FCESOptimizer opt(params, FCESConfig{}.set_lr(1e-3f));
|
|
|
|
// Simulate a training step
|
|
auto x = torch::randn({2, 10});
|
|
auto y = model->forward(x);
|
|
auto loss = y.sum();
|
|
loss.backward();
|
|
opt.step();
|
|
|
|
EXPECT_EQ(opt.step_count(), 1);
|
|
}
|
|
|
|
TEST(OptimizerTest, UpdateFitness) {
|
|
auto model = torch::nn::Linear(10, 5);
|
|
std::vector<torch::Tensor> params;
|
|
for (auto &p : model->parameters())
|
|
params.push_back(p);
|
|
|
|
FCESOptimizer opt(params);
|
|
opt.update_fitness(3.0f);
|
|
opt.update_fitness(2.5f);
|
|
// Should not crash
|
|
}
|