将上面代码改成定义两个卷基层
时间: 2024-03-24 07:38:51 浏览: 21
好的,可以参考下面的代码实现:
```python
import torch
import torch.nn as nn
class ConvBlock(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1):
super(ConvBlock, self).__init__()
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding)
self.bn = nn.BatchNorm2d(out_channels)
self.relu = nn.ReLU(inplace=True)
def forward(self, x):
x = self.conv(x)
x = self.bn(x)
x = self.relu(x)
return x
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = ConvBlock(3, 32)
self.conv2 = ConvBlock(32, 64)
self.fc = nn.Linear(64 * 28 * 28, 10)
def forward(self, x):
x = self.conv1(x)
x = self.conv2(x)
x = x.view(x.size(0), -1)
x = self.fc(x)
return x
```
这里我们定义了一个 `ConvBlock` 类,它包含了一个卷积层、一个批归一化层和一个激活函数ReLU。然后我们在 `Net` 类中定义了两个 `ConvBlock` 类型的卷积层,分别是 `self.conv1` 和 `self.conv2`。最后还有一个全连接层 `self.fc`,将卷积层输出的特征展开后接到全连接层中。