代码
#include <torch/script.h>
#include <ATen/ATen.h>
#include <torch/nn/module.h>
#include <torch/nn/modules/batchnorm.h>
#include <torch/nn/modules/conv.h>
#include <torch/nn/modules/rnn.h>
#include <torch/nn/modules/dropout.h>
#include <torch/nn/modules/embedding.h>
#include <torch/nn/modules/functional.h>
#include <torch/nn/modules/linear.h>
#include <torch/nn/modules/sequential.h>
#include <torch/optim.h>
#include <torch/types.h>
#include <torch/utils.h>
#include <cmath>
#include <cstdlib>
#include <functional>
#include <iostream>
#include <memory>
#include <random>
#include <vector>
using namespace std;
using namespace at;
using namespace torch::nn;
using namespace torch::optim;
torch::Tensor test_net(torch::Tensor input){
Sequential model(
Linear(2,8),
BatchNorm(8),
Functional(torch::sigmoid),
Linear(8,1),
Functional(torch::sigmoid));
return model->forward(input);
};
int main(int argc, const char* argv[])
{
torch::Tensor img = torch::randn({10,2});
std::cout << img.sizes() << endl;
auto out = test_net(img);
std::cout << out.sizes() << endl;
std::cout<< "ok\n";
return 1;
}
编译
make
./bin/demo
执行
[10, 2]
[10, 1]
ok