pytorch訓練簡單的CNN(visdom進行視覺化)

from visdom import Visdom #監聽資料
import torch
import torch.nn as nn
from torch .autograd import Variable #匯入自動求導機制
import torch.utils.data as Data #匯入data資料集
import torchvision #匯入torchvision資料集包，裡面包含影像翻轉等功能

#-------------------------------------------------定義超引數--------------------------------------------------------------
epoch_n = 3 #所有資料被輪到的次數
LR = 0.001 #優化演算法中所用到的學習率
BatchSize = 50

#-------------------------------------------------視覺化視窗設定-----------------------------------------------------------
viz = Visdom()#將視窗類例項化
viz.line([0.], [0], win='train_loss', opts=dict(title='train_loss'))#建立視窗並初始化（只監聽loss）
#viz.line([[0.,0.]], [0], win='train_loss', opts=dict(title='loss&acc', legend=['loss', 'acc']))#建立視窗並初始化（loss和accuracy同時監聽）

#-------------------------------------------------下載資料集torchvision.datasets.MNIST---------------------------------
train_data = torchvision.datasets.MNIST(
    root="./mnist", #root: 資料集存放位置。mnist資料夾中分為processed（包括train和test兩個資料夾）和raw兩個資料夾
    train=True, #True為訓練集，False為測試集
    transform=torchvision.transforms.ToTensor(), #將原資料變換為(0,1)區間。並Tosensor到torch能識別的tensor型別
    download=True, #True為從網路下載資料集，若已下載則不會再次下載
)

print(train_data.train_data.size()) #輸出MNIST的訓練集的大小
print(train_data.train_labels.size()) #輸出訓練集的標籤大小

#for i in range(1,4):
#    plt.imshow(train_data.train_data[i].numpy(),cmap="gray") #繪製第i張圖片
#    plt.title("%i" % train_data.train_labels[i]) #新增第i張圖片的標籤
#    plt.show()

test_data = torchvision.datasets.MNIST( #獲取測試集（下載訓練集時測試集已經下載好了，這裡獲取一下）
    root="./mnist",
    train=False,
)

#-------------------------------------------------------載入資料集Data.DataLoader-----------------------------------------
train_loader = Data.DataLoader(dataset=train_data, #這裡的Data是import torch.utils.data as Data。 dataset是已經下載好的資料集用於載入
                               batch_size=BatchSize, #載入批訓練的資料個數
                               shuffle=True) #在每個epoch重新排列的資料

test_x = Variable(torch.unsqueeze(test_data.test_data,dim=1), volatile=True).type(torch.FloatTensor) #注意這裡和訓練集的載入方式不一樣
test_y = test_data.test_labels

#--------------------------------------------------------------定義網路結構------------------------------------------------
class CNN(nn.Module): #nn是匯入的
    def __init__(self):
        super(CNN,self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(in_channels=1, out_channels=16, kernel_size=5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(in_channels=16, out_channels=32, kernel_size=5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)
        )
        self.out=nn.Linear(32*7*7,10)

    def forward(self,x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(x.size(0),-1)
        output = self.out(x)
        return output

cnn = CNN()
print(cnn) #列印建立的cnn的結構

params = list(cnn.parameters())  #列印構建的網路的引數
print(len(params))
print(params[0].size())

#----------------------------------------------------------確定優化方法和損失函式---------------------------------------------
optimizer = torch.optim.Adam(cnn.parameters(),lr=LR) #Adam優化
loss_function = nn.CrossEntropyLoss() #交叉熵損失函式

#------------------------------------------------------------進行每一輪的訓練----------------------------------------------
for epoch in range(epoch_n):

    print("Epoch {}/{}".format(epoch, epoch_n-1))
    print("-"*20) #這裡的20沒有具體的含義，只是用於分隔符的長度

    for batch, (x,y) in enumerate(train_loader):
        #前向傳播
        b_x = Variable(x)
        b_y = Variable(y)
        train_output = cnn(b_x)
        #計算訓練集的損失函式和精確度
        train_loss = loss_function(train_output, b_y)

        #反向傳播並優化
        optimizer.zero_grad() #每次反向傳播前都要清空上一次的梯度
        train_loss.backward()
        optimizer.step()

        if batch % 100 == 0:
            test_output = cnn(test_x)
            test_loss = loss_function(test_output, test_y)
            pred_y = torch.max(test_output,1)[1].data.squeeze()
            test_accuracy = sum(pred_y==test_y) / test_y.size(0)
            print("Batch:", batch, "| Train loss:%.4f"%float(train_loss), "| Test accuracy:%.4f" % test_accuracy)
            viz.line([float(train_loss)], [batch], win='train_loss', update='append') #利用visdom監聽loss
            #viz.line([[float(loss),float(accuracy)]], [step], win='train_loss', update='append') #利用visdom同時監聽loss和accuracy