用卷积神经网络优化常规BP神经网络

摘要

我是个新手，最近做项目遇到了一些问题，老是在将优秀网络模型替换其他网络的时候会出现模型大小不匹配的问题，先将自己做好的代码发出来，做个记录：

我用的minist手写字体数据集，因为我的电脑限制，我用CPU跑的，如果条件允许，可以用cuda

常规BP神经网络

BP网络结构

class BP(nn.Module):
    def __init__(self):
        super(BP, self).__init__()

        self.fc = nn.Sequential(
            nn.Linear(784, 128),
            nn.ReLU(inplace=True),
            nn.Linear(128, 64),
            nn.ReLU(inplace=True),
            nn.Linear(64, 10)
        )

    def forward(self, x):
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

整体代码

import torch
import numpy as np
from torch import nn, optim
from torch.autograd import Variable
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import matplotlib.pyplot as plt


class BP(nn.Module):
    def __init__(self):
        super(BP, self).__init__()

        self.fc = nn.Sequential(
            nn.Linear(784, 128),
            nn.ReLU(inplace=True),
            nn.Linear(128, 64),
            nn.ReLU(inplace=True),
            nn.Linear(64, 10)
        )

    def forward(self, x):
        x = x.view(x.size(0), -1)	
        x = self.fc(x)
        return x


# train
def train(model, train_loader, optimizer, criterion, device):
    model.train()
    for batch_idx, data in enumerate(train_loader):
        img, label = data

        img = img.to(device)
        label = label.to(device)

        out = model(img)
        optimizer.zero_grad()
        loss = criterion(out, label)
        loss.backward()
        optimizer.step()

        if