Pytorch学习（五）--实现神经网络实例二

'''
构建模型
'''
#构建网络
class Net(nn.Module):
    def __init__(self,in_dim,n_hidden_1,n_hidden_2,out_dim):
        super(Net,self).__init__()
        self.layer1 = nn.Sequential(nn.Linear(in_dim,n_hidden_1),nn.BatchNorm1d(n_hidden_1))
        self.layer2 = nn.Sequential(nn.Linear(n_hidden_1,n_hidden_2),nn.BatchNorm1d(n_hidden_2))
        self.layer3 = nn.Sequential(nn.Linear(n_hidden_2,out_dim))

    def forward(self,x):
        x = F.relu(self.layer1(x))
        x = F.relu(self.layer2(x))
        x = self.layer3(x)
        return x

# 实例化网络层
#检查是否有可用的GPU
# device = torch.device("cuda ：0" if torch.cuda.is_available() else "cpu")
device = torch.device("cuda:0")
model = Net(28*28 , 300,100,10)
model.to(device)

# 定义损失函数和优化器
criterion =  nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(),lr=lr,momentum=momentum)

'''
训练模型
'''
losses = []  # 存储每一个epoch的损失
acces = []   # 存储每一个epoch的损失
eval_losses = []
eval_acces = []

for epoch in range(num_epoches):
    train_loss = 0
    train_acc = 0
    model.train()  #这句话会自动调整batch_normalize和dropout值，很关键！

# 动态修改学习率
    if epoch%5 == 0:
        optimizer.param_groups[0]['lr'] *=0.1

    # 对于整个数据集，通过dataloader，每次使用batch_size个数据
    for img,label in train_loader:
        img = img.to(device)
        label = label.to(device)
        img = img.view(img.size(0),-1)  #TODO 输出一下看看
        # 前向传播
        out = model(img)  #out是64*10
        loss = criterion(out,label)
        # 反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # 记录误差 ,将每一次train_loader的loss累积起来，形成每一个epoch的误差
        train_loss += loss.item()  # pytorch中，.item()方法 是得到一个元素张量里面的元素值
                                    # 具体就是 用于将一个零维张量转换成浮点数，比如计算loss，accuracy的值
        # 计算分类的准确性
        _, pred = out.max(1)   # 返回每一行中的最大值赋给_  , 同时把这个最大值对应的列的索引赋值给pred
        num_correct = (pred == label).sum().item()
        acc = num_correct / img.shape[0]
        train_acc += acc

    losses.append(train_loss / len(train_loader))
    acces.append(train_acc / len(train_loader))
    print('epoch:{},train_loss:{:.6f},acc:{:.6f}'.format(epoch + 1, train_loss / len(train_loader),
                                                         train_acc / len(train_loader)))
    # 在测试集上检验效果
    eval_loss = 0
    eval_acc = 0

    # 讲模型改为预测模式
    model.eval() # 在评估模型时使用，固定BN 和 Dropout
    for img,label in test_loader:
        img = img.to(device)
        label = label.to(device)
        img = img.view(img.size(0), -1)
        out = model(img)
        loss = criterion(out,label)

        eval_loss +=loss.item()
        _, pred = out.max(1)
        num_correct = (pred == label).sum().item()
        acc = num_correct / img.shape[0]
        eval_acc += acc

    eval_losses.append(eval_loss / len(test_loader))
    eval_acces.append(eval_acc / len(test_loader))

    print('Test Loss:{:.6f}, Acc:{:.6f}'.format(eval_loss / len(test_loader), eval_acc / len(test_loader)))


'''
可视化训练结果

'''

plt.title('train_loss')
plt.plot(np.arange(len(losses)),losses)
plt.legend(['Train Loss'],loc= 'upper right')
plt.show()