用Pytorch訓練分類模型
阿新 • • 發佈:2019-02-07
定義 計算 tensor imp cor entropy lin epo learning
本次分類問題使用的數據集是MNIST,每個圖像的大小為\(28*28\)。
編寫代碼的步驟如下
- 載入數據集,分別為訓練集和測試集
- 讓數據集可以叠代
- 定義模型,定義損失函數,訓練模型
代碼
import torch import torch.nn as nn import torchvision.transforms as transforms import torchvision.datasets as dsets from torch.autograd import Variable '''下載訓練集和測試集''' train_dataset = dsets.MNIST(root='./datasets', train=True, transform=transforms.ToTensor(), download=True) test_dataset = dsets.MNIST(root='./datasets', train=False, transform=transforms.ToTensor()) '''讓數據集可以叠代''' batch_size = 100 n_iters = 3000 num_epochs = n_iters / (len(train_dataset) / batch_size) num_epochs = int(num_epochs) train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True) test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False) '''定義模型''' class LogisticRegressionModel(nn.Module): def __init__(self, input_dim, output_dim): super(LogisticRegressionModel, self).__init__() self.linear = nn.Linear(input_dim, output_dim) def forward(self, x): out = self.linear(x) return out '''實例化模型''' input_dim = 28*28 output_dim = 10 model = LogisticRegressionModel(input_dim, output_dim) '''定義損失計算方式''' criterion = nn.CrossEntropyLoss() learning_rate = 0.001 optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate) '''訓練次數''' iter = 0 for epoch in range(num_epochs): for i, (images, labels) in enumerate(train_loader): images = Variable(images.view(-1, 28*28)) labels = Variable(labels) #梯度置零 optimizer.zero_grad() #計算輸出 outputs = model(images) #計算損失,內部會自動softmax然後進行Crossentropy loss = criterion(outputs, labels) #反向傳播 loss.backward() #更新參數 optimizer.step() iter += 1 if iter % 500 == 0: #計算準確度 correct = 0 total = 0 for images, labels in test_loader: images = Variable(images.view(-1, 28*28)) #獲得輸出,輸出的大小為(batch_size,10) outputs = model(images) #獲得預測值,輸出的大小為(batch_size,1) _, predicted = torch.max(outputs.data, 1) #labels的size是(100,) total += labels.size(0) #返回的是預測值和標簽值相等的個數 correct += (predicted == labels).sum() accuracy = 100 * correct / total # Print Loss print('Iteration: {}. Loss: {}. Accuracy: {}'.format(iter, loss.data[0], accuracy))
輸出如下
用Pytorch訓練分類模型