Juzhen/examples/demo_rectified_infer.cu at main · anewgithubname/Juzhen · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
/**
 * @file demo_rectified_infer.cu
 * @brief using the pretrained retified flow to infer the data
 * @author Song Liu (song.liu@bristol.ac.uk)
 *
 * This is an implementation of the rectified flow
 * https://arxiv.org/abs/2209.03003
 *
    Copyright (C) 2024 Song Liu (song.liu@bristol.ac.uk)

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <https://www.gnu.org/licenses/>.

 */

// #define CPU_ONLY
#include "../cpp/juzhen.hpp"
#include "../ml/layer.hpp"
#include "../ml/dataloader.hpp"
#include "../ml/plotting.hpp"

#include <sstream>
#include <fstream>

using namespace std;
using namespace Juzhen;

#ifdef CUDA
#define FLOAT CUDAfloat
inline Matrix<CUDAfloat> randn(int m, int n) { return Matrix<CUDAfloat>::randn(m, n); }
inline Matrix<CUDAfloat> ones(int m, int n) { return Matrix<CUDAfloat>::ones(m, n); }
inline Matrix<CUDAfloat> vs(std::vector<MatrixView<CUDAfloat>> matrices) {return vstack(matrices);}
inline Matrix<CUDAfloat> hs(std::vector<MatrixView<CUDAfloat>> matrices) {return hstack(matrices);}
inline const float* getdata(const Matrix<CUDAfloat>& m) {return m.to_host().data();}
#else
#define FLOAT float
inline Matrix<float> randn(int m, int n) { return Matrix<float>::randn(m, n); }
inline Matrix<float> ones(int m, int n) { return Matrix<float>::ones(m, n); }
inline Matrix<float> vs(std::vector<MatrixView<float>> matrices) {return vstack<float>(matrices);}
inline Matrix<float> hs(std::vector<MatrixView<float>> matrices) {return hstack<float>(matrices);}
inline const float* getdata(const Matrix<float>& m) {return m.data();}
#endif
#define MatrixI Matrix<int>


void PrintSeparationLine();

auto sample_X0(int n, int d)
{
    return randn(d, n);
}

int compute() {
    // spdlog::set_level(spdlog::level::debug);
#ifdef CUDA
    GPUSampler sampler(1);
#endif

    Profiler p("comp");
    using namespace Juzhen;
    std::string base = PROJECT_DIR;

    int d = 1;
    int n = 6500;

    auto X0 = sample_X0(n, d); // reference data
    plot_histogram(getdata(X0), X0.num_row() * X0.num_col(), 23);

    const size_t batchsize = 1;

    // create a neural network
    // define layers
    ReluLayer<FLOAT> L0(1003, d + 1, batchsize),
                     L1(1003, 1003, batchsize),
                     L2(1003, 1003, batchsize),
                     L3(1003, 1003, batchsize);
    LinearLayer<FLOAT> L10(d, 1003, batchsize);

    // nns are linked lists containing layers
    list<Layer<FLOAT> *> trainnn({&L10, &L3, &L2, &L1, &L0});

    loadweights(trainnn, base+"/res/net.weights");

    PrintSeparationLine();

    X0 = sample_X0(n, d); // reference data
    auto Zt_trajectory = euler_integration(X0, trainnn, 100);

    for (auto Zt : Zt_trajectory) {
        plot_histogram(getdata(Zt), Zt.num_row() * Zt.num_col(), 23);
    }

    return 0;
}