构建卷积网络

#include<torch/torch.h>
#include<torch/script.h>
#include<iostream>

using std::cout; using std::endl;

class LinearBnReluImpl : public torch::nn::Module
{
private:
	torch::nn::Linear ln{ nullptr };
	torch::nn::BatchNorm1d bn{ nullptr };

public:
	LinearBnReluImpl(int input_features, int out_features);
	torch::Tensor forward(torch::Tensor x);
};
TORCH_MODULE(LinearBnRelu);

inline torch::nn::Conv2dOptions conv_options(
	int64_t in_planes, int64_t out_planes, int64_t kernel_size,
	int64_t stride = 1, int64_t padding = 0, bool with_bias = false
)
{
	torch::nn::Conv2dOptions conv_options = torch::nn::Conv2dOptions(in_planes, out_planes, kernel_size);
	conv_options.stride(stride);
	conv_options.padding(padding);
	conv_options.bias(with_bias);

	return conv_options;
}

class ConvReluBnImpl : public torch::nn::Module
{
private:
	torch::nn::Conv2d conv{ nullptr };
	torch::nn::BatchNorm2d bn{ nullptr };
	
public:
	ConvReluBnImpl(int input_channel, int output_channel, int kernel_size, int stride, int padding=1);
	torch::Tensor forward(torch::Tensor x);
};
TORCH_MODULE(ConvReluBn);


class MLP : public torch::nn::Module
{
private:
	int mid_features[3] = { 32, 64, 128 };
	LinearBnRelu ln1{ nullptr };
	LinearBnRelu ln2{ nullptr };
	LinearBnRelu ln3{ nullptr };
	torch::nn::Linear out_ln{ nullptr };

public:
	MLP(int in_features, int out_features);
	torch::Tensor forward(torch::Tensor x);
};

class plainCNN : public torch::nn::Module
{
private:
	int mid_channels[3]{ 32,64,128 };
	ConvReluBn conv1{ nullptr };
	ConvReluBn down1{ nullptr };
	ConvReluBn conv2{ nullptr };
	ConvReluBn down2{ nullptr };
	ConvReluBn conv3{ nullptr };
	ConvReluBn down3{ nullptr };
	torch::nn::Conv2d out_conv{ nullptr };

public:
	plainCNN(int in_channels, int out_channels);
	torch::Tensor forward(torch::Tensor x);
};


int main()
{
	plainCNN c(3, 2);

	auto x = torch::rand({ 1,3,224,224 }, torch::kFloat);
	//cout << x.sizes() << endl;
	
	auto a = c.forward(x);
	cout <<"[in Main]: "<< a.sizes() << endl;

	return 0;
}

LinearBnReluImpl::LinearBnReluImpl(int input_features, int out_features)
{
	ln = register_module("ln", torch::nn::Linear(torch::nn::LinearOptions(input_features, out_features)));
	bn = register_module("bn", torch::nn::BatchNorm1d(out_features));
}

torch::Tensor LinearBnReluImpl::forward(torch::Tensor x)
{
	x = torch::relu(ln->forward(x));
	x = bn(x);
	return x;
}

ConvReluBnImpl::ConvReluBnImpl(int input_channel, int output_channel, int kernel_size, int stride, int padding)
{
	conv = register_module("conv", torch::nn::Conv2d(conv_options(input_channel, output_channel, kernel_size, stride, padding)));
	bn = register_module("bn", torch::nn::BatchNorm2d(output_channel));
}

torch::Tensor ConvReluBnImpl::forward(torch::Tensor x)
{
	x = torch::relu(conv->forward(x));
	x = bn(x);
	return x;
}

MLP::MLP(int in_features, int out_features)
{
	ln1 = LinearBnRelu(in_features, mid_features[0]);
	ln2 = LinearBnRelu(mid_features[0], mid_features[1]);
	ln3 = LinearBnRelu(mid_features[1], mid_features[2]);
	out_ln = torch::nn::Linear(mid_features[2], out_features);

	ln1 = register_module("ln1", ln1);
	ln2 = register_module("ln2", ln2);
	ln3 = register_module("ln3", ln3);
	out_ln = register_module("out_ln", out_ln);
}

torch::Tensor MLP::forward(torch::Tensor x)
{
	x = ln1->forward(x);
	x = ln2->forward(x);
	x = ln3->forward(x);
	x = out_ln->forward(x);
	return x;
}

plainCNN::plainCNN(int in_channels, int out_channels)
{
	conv1 = ConvReluBn(in_channels, mid_channels[0], 3, 1);
	down1 = ConvReluBn(mid_channels[0], mid_channels[0], 3, 2);
	conv2 = ConvReluBn(mid_channels[0], mid_channels[1], 3,1);
	down2 = ConvReluBn(mid_channels[1], mid_channels[1], 3, 2);
	conv3 = ConvReluBn(mid_channels[1], mid_channels[2], 3,1);
	down3 = ConvReluBn(mid_channels[2], mid_channels[2], 3, 2);
	out_conv = torch::nn::Conv2d(conv_options(mid_channels[2], out_channels, 3));

	conv1 = register_module("conv1", conv1);
	down1 = register_module("down1", down1);
	conv2 = register_module("conv2", conv2);
	down2 = register_module("down2", down2);
	conv3 = register_module("conv3", conv3);
	down3 = register_module("down3", down3);
	out_conv = register_module("out_conv", out_conv);
}

torch::Tensor plainCNN::forward(torch::Tensor x)
{
	x = conv1->forward(x);
	cout << x.sizes() << endl;
	x = down1->forward(x);
	cout << x.sizes() << endl;
	x = conv2->forward(x);
	cout << x.sizes() << endl;
	x = down2->forward(x);
	cout << x.sizes() << endl;
	x = conv3->forward(x);
	cout << x.sizes() << endl;
	x = down3->forward(x);
	cout << x.sizes() << endl;
	x = out_conv->forward(x);
	cout << x.sizes() << endl;
	return x;
}

结果