帮我写一份cnn的网络代码
时间: 2023-07-16 18:12:38 浏览: 33
好的,以下是一个简单的CNN网络的代码:
```python
import torch
import torch.nn as nn
class CNN(nn.Module):
def __init__(self):
super(CNN, self).__init__()
self.conv1 = nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3, stride=1, padding=1)
self.relu1 = nn.ReLU()
self.pool1 = nn.MaxPool2d(kernel_size=2)
self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=1, padding=1)
self.relu2 = nn.ReLU()
self.pool2 = nn.MaxPool2d(kernel_size=2)
self.fc1 = nn.Linear(64 * 8 * 8, 512)
self.relu3 = nn.ReLU()
self.fc2 = nn.Linear(512, 10)
def forward(self, x):
x = self.conv1(x)
x = self.relu1(x)
x = self.pool1(x)
x = self.conv2(x)
x = self.relu2(x)
x = self.pool2(x)
x = x.view(-1, 64 * 8 * 8)
x = self.fc1(x)
x = self.relu3(x)
x = self.fc2(x)
return x
```
这个网络包含两个卷积层和两个全连接层。输入图像的大小应该为 $3 \times 32 \times 32$,输出为 $10$ 个类别的概率分布。你可以根据自己的数据集和任务进行修改。