如何使用Pytorch迅速寫一個Mnist數據分類器

如何使用Pytorch迅速寫一個Mnist數據分類器

一段時間沒有更新博文,想着也該寫兩篇文章玩玩了。而從一個簡單的例子做爲開端是一個比較不錯的選擇。本文章會手把手地教讀者構建一個簡單的Mnist(Fashion-Mnist同理)的分類器,而且會使用相對完整的Pytorch訓練框架,所以對於初學者來講應該會是一個方便入門且便於閱讀的文章。本文的代碼來源於我剛學Pytorch時的小項目,可能在形式上會有引用一些github上的小代碼。同時文風可能會和我以前看的一些外國博客有點相近。python

本文適用對象: 剛入門的Pytorch新手,想要用Pytorch來完成做業的魚乾。

那麼就開始coding吧。git

首先,你須要安裝好Python 3+,Pytorch 1.0+,我我的使用的是Pytorch1.4,我想1.0以上的版本均可以使用。github

而後在想要的位置,新建一個main.py的文件,而後就能夠開始敲鍵盤了。網絡

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import  datasets
from torchvision import transforms
import torch.utils.data

import argparse

第一步天然是導入相應的包。前面的都是Pytorch的包,最後一句導入的argparse便於用來修改訓練的參數,這在Pytorch復現深度學習模型時很是常見。框架

model_names = ['Net','Net1']

parser = argparse.ArgumentParser(description='PyTorch Mnist Training')
parser.add_argument('-a', '--arch', metavar='ARCH', default='Net', choices=model_names,
                    help='model archtecture: ' + '|'.join(model_names) + '(default:Net)')
parser.add_argument('--epochs', default=5, type=int, metavar='N', help='number of total epochs')
parser.add_argument('--momentum', default=0.9, type=float, metavar='M', help='momentum')
parser.add_argument('-b', '--batchsize', default=32, type=int, metavar='N', help='mini-batch size')
parser.add_argument('--lr', '--learning-rate', default=1e-2, type=float, metavar='LR', help='initial learning rata',
                    dest='lr')
args=vars(parser.parse_args())

第一行的model_names是一個list,用來存儲咱們以後會實現的兩種網絡結構的名字。而後我定義了一個argparse的對象,關於argparse能夠自尋一些教程觀看,大概只須要知道能夠從指令行輸入參數便可。在parser中又定義了arch(使用的網絡),epochs(迭代輪次),momentum(梯度動量大小),batchsize(一次送入的圖片量大小),learning-rate(學習率)參數。以前的model_name也正是用在arch參數中,限定了網絡框架將會今後兩者中選擇其一。ide

def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    #parameter
    batch_size = args["batchsize"]
    lr = args["lr"]
    momentum = args["momentum"]
    num_epochs = args["epochs"]

主函數中,先定義cuda對象,便於使用gpu並行運算。在#parameter中,咱們把一些從命令行中得到的參數引入到相應的變量中,以便後續書寫。函數

#prepare the dataset

    mnist_data = datasets.MNIST("./mnist_data",train=True, download=True,
                                transform=transforms.Compose([
                                    transforms.ToTensor(),
                                    transforms.Normalize(mean=(0.13,),std=(0.31,))
                                ]))
    '''
    mnist_data = datasets.FashionMNIST("./fashion_mnist_data", train=True, download=True,
                                transform=transforms.Compose([
                                    transforms.ToTensor(),
                                    transforms.Normalize(mean=(,), std=(,))
                                ]))
    '''
    train_loader = torch.utils.data.DataLoader(
        mnist_data,batch_size=batch_size,shuffle=True, num_workers=1,pin_memory=True
    )
    test_loader = torch.utils.data.DataLoader(
        mnist_data,batch_size=batch_size,shuffle=True, num_workers=1,pin_memory=True
    )

以後將引入Pytorch中datasets包自帶的MNIST集,download參數設置爲True,以便於本地沒有Mnist數據集時直接下載,以後會在當前目錄下建立一個mnist_data的文件夾以存放數據,。transform中的transforms.ToTensor()是用於將圖片形式的數據轉換成tensor類型,而transforms.Normalize(mean=(0.13,),std=(0.31,))則是將tensor類型的數據進行歸一化,這裏的0.13和0.31能夠直接使用。若是你想要使用註釋中的FashionMNIST數據集則須要使用的是註釋中的內容,固然,mean和std須要另外求解。學習

以後,定義train_loadertest_loader,將數據集做爲可迭代的對象使用。shuffle=True以實現亂序讀取數據,通常都會這麼設置。num_workerspin_memory都會影響到數據讀取速度,前者是會在讀取時建立多少個進程,後者是影響到數據讀入到GPU中,通常來講,對於這個項目前者設置爲1已經夠用,後者設置爲True和False都不影響。在更大型的項目中,若是設備較好,前者能夠設置大一些。測試

model = Net1().to(device) if args["arch"]=='Net1' else Net().to(device)
    optimizer = torch.optim.SGD(model.parameters(),lr=lr,momentum=momentum)
    criteon = nn.CrossEntropyLoss().to(device)

第一行是model實例化,而且會根據args["arch"]選擇是用Net仍是Net1to(device)會將model放置於device上運行。第二行定義了一個優化器,使用的是SGD,而且放入model的參數、學習率和動量大小。criteon定義損失函數,這邊使用的是交叉熵函數,這一損失函數在分類問題中十分經常使用。優化

#train
    for epoch in range(num_epochs):
        train(model,device,train_loader,optimizer,epoch,criteon)
        test(model,device,test_loader,criteon)


    torch.save(model.state_dict(), "mnist_{}.pth".format(num_epochs))

這就是訓練過程,在其中又使用了traintest兩個函數(下面會說),根據num_epochs數目進行循環。循環結束後,torch.save將會把模型的參數model.state_dict()mnist_{}.pth的形式存放到當前文件夾下。

def train(model,device,train_loader,optimizer,epoch,criteon):
    class_name = model.__class__.__name__

    model.train()
    loss = 0
    for idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        pred = model(data)
        if class_name == 'Net':
            loss = F.nll_loss(pred, target)
        elif class_name == 'Net1':
            loss = criteon(pred, target)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if idx % 100 == 0:
            print("train epoch: {}, iteration: {}, loss: {}".format(epoch, idx, loss.item()))

這裏定義了train函數的訓練過程。class_name中存放了當前使用的模型名字。 model.train()開啓訓練模式。在for idx, (data, target) in enumerate(train_loader):取出當前數據集的idx,data和種類target。循環中,先把data和target放置於device上,pred = model(data)會進行一次前傳,得到相應數據的預測種類pred

對不一樣的模型,我採用了不一樣定義損失函數的方式,這裏須要結合下面的模型結構來看。optimizer.zero_grad()會將上輪累計的梯度清空,以後loss.backward()梯度反向傳播,利用optimizer.step()更新參數。而當if idx % 100 == 0:也就是迭代的數據批次到達100的倍數了,就會輸出相關信息。

def test(model,device,test_loader,criteon):
    class_name = model.__class__.__name__
    model.eval()
    total_loss = 0 #caculate total loss
    correct = 0
    with torch.no_grad():
        for idx, (data, target) in enumerate(test_loader):
            data, target = data.to(device), target.to(device)
            pred = model(data)
            if class_name == 'Net':
                total_loss += F.nll_loss(pred, target,reduction="sum").item()
            elif class_name == 'Net1':
                total_loss += criteon(pred, target).item()
            correct += pred.argmax(dim=1).eq(target).sum().item()

    total_loss /= len(test_loader.dataset)
    acc = correct/len(test_loader.dataset)
    print("Test loss: {}, Accuracy: {}%".format(total_loss,acc*100))

test函數整體結構相似,model.eval()將會把模型調整測試模式,with torch.no_grad():來聲明測試模式下不須要積累梯度信息。correct += pred.argmax(dim=1).eq(target).sum().item()則是會計算出預測對了的數目,以後經過total_loss計算總偏差和acc計算準確率。

class Net(nn.Module):
    def __init__(self):
        super(Net,self).__init__()
        self.conv1 = nn.Conv2d(1, 20, kernel_size=5,stride=1)
        self.conv2 = nn.Conv2d(20,50,kernel_size=5,stride=1)
        self.fc1 = nn.Linear(4*4*50, 500)
        self.fc2 = nn.Linear(500, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.max_pool2d(x,2,2)
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x,2,2)
        x = x.view(-1,4*4*50)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        x = F.log_softmax(x,dim=1)
        return x

Net不過是一個具備兩個卷積層和兩個線性全連層的網絡。self.conv1 = nn.Conv2d(1, 20, kernel_size=5,stride=1)表示conv1是一個接受1個channel的tensor輸出20個channel的tensor,且卷積大小爲5,步長爲1的卷積層。self.fc1 = nn.Linear(4*4*50, 500)則是接收一個4 * 4 * 50長的一維tensor而且輸出長爲500的一維tensor。

前傳函數forward中,x做爲輸入的數據,輸入後會經過conv1->relu->pooling->conv2->relu->pooling->view將多維tensor轉化成一維tensor->fc1->relu->fc2->log_softmax來得到最終的x的值。這裏就須要提train和test函數中的if和elif語句了。使用的時Net時,loss = F.nll_loss(pred, target),這是由於log_softmax以後使用nll_loss和直接使用 nn.CrossEntropyLoss()是等效的,所以:

class Net1(nn.Module):
    def __init__(self):
        super(Net1,self).__init__()
        self.conv_unit=nn.Sequential(
            nn.Conv2d(1, 20, kernel_size=5,stride=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2,stride=2),
            nn.Conv2d(20,50,kernel_size=5,stride=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        self.fc_unit=nn.Sequential(
            nn.Linear(4*4*50, 500),
            nn.ReLU(),
            nn.Linear(500, 10)
        )


    def forward(self, x):
        x = self.conv_unit(x)
        x = x.view(-1,4*4*50)
        x = self.fc_unit(x)
        return x

Net1中最後並無使用log_softmax,是由於直接在train的過程當中,使用了nn.CrossEntropyLoss()。此外,Net1和Net不一樣的地方也就是在結構中使用了nn.Sequential()來單元化卷積層和全連層。

if __name__ == '__main__':
    main()

以後就可使用了!

在命令行中使用:

$ python main.py

就會按照默認的參數訓練一個Mnist分類器了。

第三輪的效果:

train epoch: 2, iteration: 1300, loss: 0.010509848594665527
train epoch: 2, iteration: 1400, loss: 0.0020529627799987793
train epoch: 2, iteration: 1500, loss: 0.0027058571577072144
train epoch: 2, iteration: 1600, loss: 0.010049819946289062
train epoch: 2, iteration: 1700, loss: 0.0352507084608078
train epoch: 2, iteration: 1800, loss: 0.009431719779968262
Test loss: 0.01797709318200747, Accuracy: 99.42833333333333%

若是但願查看參數列表,則能夠在命令行使用:

$ python main.py -h

就會出現:

usage: main.py [-h] [-a ARCH] [--epochs N] [--momentum M] [-b N] [--lr LR]

PyTorch Mnist Training

optional arguments:
  -h, --help            show this help message and exit
  -a ARCH, --arch ARCH  model archtecture: Net|Net1(default:Net)
  --epochs N            number of total epochs
  --momentum M          momentum
  -b N, --batchsize N   mini-batch size
  --lr LR, --learning-rate LR
                        initial learning rata

因而若是想要使用Net1,lr爲0.001的方式訓練,就能夠按照這樣:

$ python main.py -a Net1 --lr 0.001

第三輪結果:

train epoch: 2, iteration: 1200, loss: 0.03096039593219757
train epoch: 2, iteration: 1300, loss: 0.060124486684799194
train epoch: 2, iteration: 1400, loss: 0.08865253627300262
train epoch: 2, iteration: 1500, loss: 0.13717596232891083
train epoch: 2, iteration: 1600, loss: 0.003894627094268799
train epoch: 2, iteration: 1700, loss: 0.06881710141897202
train epoch: 2, iteration: 1800, loss: 0.03184908628463745
Test loss: 0.0013615453808257978, Accuracy: 98.69666666666667%

至此,你得到了一個Mnist訓練器的訓練方法。

相關文章
相關標籤/搜索