Flow Matching

1
// Compile with: g++ -O2 -std=c++17 cfm_train.cpp -I${LIBTORCH}/include -I${LIBTORCH}/include/torch/csrc/api/include -L${LIBTORCH}/lib -ltorch -lc10 -o cfm_train -Wl,-rpath,${LIBTORCH}/lib
2
// This example trains v_θ(x, t, c) on a simple 2D Gaussian mixture with class labels.
3
#include <torch/torch.h>
4
#include <iostream>
5
#include <vector>
6
#include <random>
7
8
// Simple MLP for v_θ(x, t, c)
9
struct VelocityNetImpl : torch::nn::Module {
10
    torch::nn::Sequential net;
11
    VelocityNetImpl(int in_dim, int hidden, int out_dim) {
12
        net = torch::nn::Sequential(
13
            torch::nn::Linear(in_dim, hidden),
14
            torch::nn::SiLU(),
15
            torch::nn::Linear(hidden, hidden),
16
            torch::nn::SiLU(),
17
            torch::nn::Linear(hidden, out_dim)
18
        );
19
        register_module("net", net);
20
    }
21
    torch::Tensor forward(torch::Tensor x) {
22
        return net->forward(x);
23
    }
24
};
25
TORCH_MODULE(VelocityNet);
26
27
// One-hot encode class labels 0..K-1
28
torch::Tensor one_hot(torch::Tensor labels, int K) {
29
    auto B = labels.size(0);
30
    auto oh = torch::zeros({B, K}, labels.options().dtype(torch::kFloat32));
31
    oh.scatter_(1, labels.view({B, 1}), 1.0f);
32
    return oh;
33
}
34
35
// Sample a batch from a 2D class-conditional Gaussian mixture
36
// K classes with means on a circle; small isotropic noise
37
std::pair<torch::Tensor, torch::Tensor> sample_data(int B, int K, float radius, float stddev, torch::Device device) {
38
    auto labels = torch::randint(0, K, {B}, torch::TensorOptions().dtype(torch::kLong).device(device));
39
    auto theta = (labels.to(torch::kFloat32) * (2.0f * M_PI / K));
40
    auto means = torch::stack({radius * torch::cos(theta), radius * torch::sin(theta)}, 1); // Bx2
41
    auto eps = stddev * torch::randn({B, 2}, torch::TensorOptions().dtype(torch::kFloat32).device(device));
42
    auto x = means + eps;
43
    return {x, labels};
44
}
45
46
int main() {
47
    torch::manual_seed(42);
48
    torch::Device device(torch::kCPU);
49
50
    const int d = 2;         // data dimension
51
    const int K = 4;         // number of classes (conditions)
52
    const int in_dim = d + 1 + K; // x (2) + time t (1) + one-hot c (K)
53
    const int hidden = 128;
54
    const int out_dim = d;   // velocity vector in R^2
55
56
    VelocityNet model(in_dim, hidden, out_dim);
57
    model->to(device);
58
59
    torch::optim::Adam opt(model->parameters(), torch::optim::AdamOptions(1e-3));
60
61
    const int B = 512;       // batch size
62
    const int iters = 2000;  // training iterations
63
    const float radius = 2.5f;
64
    const float data_std = 0.15f;
65
66
    for (int it = 1; it <= iters; ++it) {
67
        model->train();
68
        // 1) Sample conditional data x1 ~ π1(·|c)
69
        auto [x1, y] = sample_data(B, K, radius, data_std, device); // x1: Bx2, y: B
70
        // 2) Sample base noise x0 ~ N(0, I)
71
        auto x0 = torch::randn_like(x1);
72
        // 3) Sample time t ~ Uniform(0, 1)
73
        auto t = torch::rand({B, 1}, torch::TensorOptions().dtype(torch::kFloat32).device(device));
74
        // 4) Build straight path x_t and target velocity u_t
75
        auto x_t = (1.0f - t) * x0 + t * x1;        // Bx2
76
        auto u_t = (x1 - x0);                       // Bx2 (constant over t for straight paths)
77
        // 5) Encode inputs: [x_t, t, one_hot(c)]
78
        auto c_oh = one_hot(y, K);
79
        auto input = torch::cat({x_t, t, c_oh}, 1); // B x (2+1+K)
80
        // 6) Predict velocity and compute MSE loss
81
        auto v_pred = model->forward(input);
82
        auto loss = torch::mse_loss(v_pred, u_t);
83
        // 7) Optimize
84
        opt.zero_grad();
85
        loss.backward();
86
        opt.step();
87
88
        if (it % 200 == 0) {
89
            std::cout << "Iter " << it << ", loss = " << loss.item<float>() << std::endl;
90
        }
91
    }
92
93
    // Save the model
94
    torch::save(model, "cfm_velocity.pt");
95
    std::cout << "Training complete. Model saved to cfm_velocity.pt\n";
96
    return 0;
97
}
98

1
// Compile with: g++ -O2 -std=c++17 cfm_sample.cpp -I${LIBTORCH}/include -I${LIBTORCH}/include/torch/csrc/api/include -L${LIBTORCH}/lib -ltorch -lc10 -o cfm_sample -Wl,-rpath,${LIBTORCH}/lib
2
#include <torch/torch.h>
3
#include <iostream>
4
#include <vector>
5
6
struct VelocityNetImpl : torch::nn::Module {
7
    torch::nn::Sequential net;
8
    VelocityNetImpl(int in_dim, int hidden, int out_dim) {
9
        net = torch::nn::Sequential(
10
            torch::nn::Linear(in_dim, hidden),
11
            torch::nn::SiLU(),
12
            torch::nn::Linear(hidden, hidden),
13
            torch::nn::SiLU(),
14
            torch::nn::Linear(hidden, out_dim)
15
        );
16
        register_module("net", net);
17
    }
18
    torch::Tensor forward(torch::Tensor x) { return net->forward(x); }
19
};
20
TORCH_MODULE(VelocityNet);
21
22
// One-hot encode labels 0..K-1
23
torch::Tensor one_hot(torch::Tensor labels, int K) {
24
    auto B = labels.size(0);
25
    auto oh = torch::zeros({B, K}, labels.options().dtype(torch::kFloat32));
26
    oh.scatter_(1, labels.view({B, 1}), 1.0f);
27
    return oh;
28
}
29
30
// Evaluate v_θ(x, t, c) given x (Bxd), scalar t, and labels y
31
torch::Tensor eval_velocity(VelocityNet &model, torch::Tensor x, float t_scalar, torch::Tensor y, int K) {
32
    auto B = x.size(0);
33
    auto t = torch::full({B, 1}, t_scalar, x.options());
34
    auto c_oh = one_hot(y, K).to(x.device());
35
    auto inp = torch::cat({x, t, c_oh}, 1);
36
    return model->forward(inp);
37
}
38
39
// One RK4 step: x_{n+1} = x_n + (h/6)(k1 + 2k2 + 2k3 + k4)
40
// where k1 = v(t_n, x_n), k2 = v(t_n + h/2, x_n + h k1/2), etc.
41
torch::Tensor rk4_step(VelocityNet &model, torch::Tensor x, float t, float h, torch::Tensor y, int K) {
42
    auto k1 = eval_velocity(model, x, t, y, K);
43
    auto k2 = eval_velocity(model, x + 0.5f * h * k1, t + 0.5f * h, y, K);
44
    auto k3 = eval_velocity(model, x + 0.5f * h * k2, t + 0.5f * h, y, K);
45
    auto k4 = eval_velocity(model, x + h * k3, t + h, y, K);
46
    return x + (h / 6.0f) * (k1 + 2.0f * k2 + 2.0f * k3 + k4);
47
}
48
49
int main() {
50
    torch::manual_seed(0);
51
    torch::Device device(torch::kCPU);
52
53
    const int d = 2;
54
    const int K = 4;
55
    const int in_dim = d + 1 + K;
56
    const int hidden = 128;
57
    const int out_dim = d;
58
59
    // Load trained model (must match architecture used during training)
60
    VelocityNet model(in_dim, hidden, out_dim);
61
    try {
62
        torch::load(model, "cfm_velocity.pt");
63
    } catch (const c10::Error& e) {
64
        std::cerr << "Failed to load model cfm_velocity.pt. Did you run training first?\n";
65
        return 1;
66
    }
67
    model->to(device);
68
    model->eval();
69
70
    const int B = 16;          // number of samples to generate
71
    auto y = torch::randint(0, K, {B}, torch::TensorOptions().dtype(torch::kLong).device(device));
72
73
    // Start from base noise x0 ~ N(0, I)
74
    auto x = torch::randn({B, d}, torch::TensorOptions().dtype(torch::kFloat32).device(device));
75
76
    // Integrate from t=0 to t=1 with RK4
77
    const int steps = 64;      // increase for higher fidelity
78
    const float h = 1.0f / steps;
79
    float t = 0.0f;
80
81
    torch::NoGradGuard no_grad;
82
    for (int s = 0; s < steps; ++s) {
83
        x = rk4_step(model, x, t, h, y, K);
84
        t += h;
85
    }
86
87
    // Print generated samples and labels
88
    auto x_cpu = x.cpu();
89
    auto y_cpu = y.cpu();
90
    auto x_a = x_cpu.accessor<float, 2>();
91
    auto y_a = y_cpu.accessor<long, 1>();
92
    for (int i = 0; i < B; ++i) {
93
        std::cout << "sample " << i << ": label=" << y_a[i]
94
                  << ", x=[" << x_a[i][0] << ", " << x_a[i][1] << "]\n";
95
    }
96
97
    return 0;
98
}
99