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 ?# 分類的類別個數
num_epochs = 10 ?# 迭代次數
batch_size = 50 ?# 批量訓練個數
learning_rate = 0.01 ?# 學習率


# 下載訓練數據和測試數據
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)

# 創建邏輯回歸類模型 ?(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

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


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

?? ??? ?# 獲得模型的訓練數據結果
? ? ? ? outputs = model(images)
?? ??? ?
?? ??? ?# 計算損失函數用于計算梯度
? ? ? ? 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()))

# 訓練好的模型預測測試數據集
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()函數中的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)