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

input_size = 784 ?# 輸入到邏輯回歸模型中的輸入大小
num_classes = 10 ?# 分類的類別個數(shù)
num_epochs = 10 ?# 迭代次數(shù)
batch_size = 50 ?# 批量訓練個數(shù)
learning_rate = 0.01 ?# 學習率


# 下載訓練數(shù)據(jù)和測試數(shù)據(jù)
train_dataset = dataset.MNIST(root='./data', train=True, transform=transforms.ToTensor(), download=True)
test_dataset = dataset.MNIST(root='./data',train=False, transform=transforms.ToTensor)

# 使用DataLoader形成批處理文件
train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False)

# 創(chuàng)建邏輯回歸類模型 ?(sigmoid(wx+b))
class LogisticRegression(nn.Module):
? ? def __init__(self,input_size,num_classes):
? ? ? ? super(LogisticRegression,self).__init__()
? ? ? ? self.linear = nn.Linear(input_size,num_classes)
? ? ? ? self.sigmoid = nn.Sigmoid()

? ? def forward(self, x):
? ? ? ? out = self.linear(x)
? ? ? ? out = self.sigmoid(out)
? ? ? ? return out

# 設定模型參數(shù)
model = LogisticRegression(input_size, num_classes)
# 定義損失函數(shù)，分類任務，使用交叉熵
criterion = nn.CrossEntropyLoss()
# 優(yōu)化算法，隨機梯度下降，lr為學習率，獲得模型需要更新的參數(shù)值
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)


# 使用訓練數(shù)據(jù)訓練模型
for epoch in range(num_epochs):
?? ?# 批量數(shù)據(jù)進行模型訓練
? ? for i, (images, labels) in enumerate(train_loader):
? ? ?? ?# 需要將數(shù)據(jù)轉換為張量Variable
? ? ? ? images = Variable(images.view(-1, 28*28))
? ? ? ? labels = Variable(labels)
?? ??? ?
?? ??? ?# 梯度更新前需要進行梯度清零
? ? ? ? optimizer.zero_grad()

?? ??? ?# 獲得模型的訓練數(shù)據(jù)結果
? ? ? ? outputs = model(images)
?? ??? ?
?? ??? ?# 計算損失函數(shù)用于計算梯度
? ? ? ? loss = criterion(outputs, labels)

?? ??? ?# 計算梯度
? ? ? ? loss.backward()
?? ?
?? ??? ?# 進行梯度更新
? ? ? ? optimizer.step()

?? ??? ?# 每隔一段時間輸出一個訓練結果
? ? ? ? if (i+1) % 100 == 0:
? ? ? ? ? ? print('Epoch:[%d %d], Step:[%d/%d], Loss: %.4f' % (epoch+1,num_epochs,i+1,len(train_dataset)//batch_size,loss.item()))

# 訓練好的模型預測測試數(shù)據(jù)集
correct = 0
total = 0
for images, labels in test_loader:
? ? images = Variable(images.view(-1, 28*28)) ?# 形式為(batch_size,28*28)
? ? outputs = model(images)
? ? _,predicts = torch.max(outputs.data,1) ?# _輸出的是最大概率的值，predicts輸出的是最大概率值所在位置，max()函數(shù)中的1表示維度，意思是計算某一行的最大值
? ? total += labels.size(0)
? ? correct += (predicts==labels).sum()

print('Accuracy of the model on the 10000 test images: %d %%' % (100 * correct / total))

train_dataset = dataset.MNIST(root='./data', train=True, transform=transforms.ToTensor(), download=True)
test_dataset = dataset.MNIST(root='./data',train=False, transform=transforms.ToTensor)

criterion = nn.CrossEntropyLoss()
loss = criterion(outputs, labels)

loss = nn.CrossEntropyLoss()
loss = loss(outputs, labels)