修改上面的模型使其一致
时间: 2023-10-09 19:13:41 浏览: 32
好的,根据错误信息,我们可以手动修改LeNet模型的结构,使其与预训练模型的结构一致。修改后的模型代码如下:
```python
import torch.nn as nn
import torch.nn.functional as F
class LeNet(nn.Module):
def __init__(self):
super(LeNet, self).__init__()
self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
self.pool1 = nn.MaxPool2d(kernel_size=2)
self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
self.pool2 = nn.MaxPool2d(kernel_size=2)
self.fc1 = nn.Linear(320, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10)
def forward(self, x):
x = F.relu(self.conv1(x))
x = self.pool1(x)
x = F.relu(self.conv2(x))
x = self.pool2(x)
x = x.view(-1, 320)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x
```
这个修改后的模型与预训练模型的结构完全一致,可以正常加载预训练模型的权重了。