写一个简单的cnn网络
时间: 2023-08-04 16:06:20 浏览: 36
以下是一个简单的CNN网络的示例代码,它包括一个卷积层和一个全连接层:
```
#include <iostream>
#include <torch/torch.h>
using namespace torch;
struct Net : nn::Module {
Net() {
// 定义卷积层
conv1 = register_module("conv1", nn::Conv2d(nn::Conv2dOptions(1, 10, /*kernel_size=*/5)));
// 定义全连接层
fc1 = register_module("fc1", nn::Linear(10 * 20 * 20, 3));
}
Tensor forward(Tensor x) {
// 进行卷积运算
x = conv1(x);
// 计算激活函数
x = torch::relu(x);
// 进行池化运算
x = torch::max_pool2d(x, /*kernel_size=*/2);
// 将输出展平为一维向量
x = torch::flatten(x, /*start_dim=*/1);
// 进行全连接运算
x = fc1(x);
// 计算概率分布
x = torch::log_softmax(x, /*dim=*/1);
return x;
}
nn::Conv2d conv1;
nn::Linear fc1;
};
int main() {
// 创建网络实例
Net net;
// 定义损失函数
nn::CrossEntropyLoss criterion;
// 定义优化器
torch::optim::SGD optimizer(net.parameters(), /*lr=*/0.01);
// 构造示例输入和输出数据
Tensor input = torch::randn({1, 1, 28, 28});
Tensor target = torch::tensor({2});
// 训练网络
for (size_t epoch = 1; epoch <= 10; ++epoch) {
// 将梯度清零
optimizer.zero_grad();
// 进行前向传播
Tensor output = net.forward(input);
// 计算损失
Tensor loss = criterion(output, target);
// 计算梯度
loss.backward();
// 更新权重
optimizer.step();
// 输出当前的损失值
std::cout << "Epoch: " << epoch << ", Loss: " << loss.item<float>() << std::endl;
}
return 0;
}
```
这个网络输入28x28的灰度图像,包括一个卷积层和一个全连接层。卷积层包括10个5x5的卷积核,可以提取局部特征。全连接层将卷积层的输出展平为一维向量,并输出三个类别的概率分布。网络使用交叉熵损失函数和随机梯度下降优化器进行训练。