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 ?# 批量訓(xùn)練個數(shù)
learning_rate = 0.01 ?# 學(xué)習(xí)率


# 下載訓(xùn)練數(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è)定模型參數(shù)
model = LogisticRegression(input_size, num_classes)
# 定義損失函數(shù)，分類任務(wù)，使用交叉熵
criterion = nn.CrossEntropyLoss()
# 優(yōu)化算法，隨機(jī)梯度下降，lr為學(xué)習(xí)率，獲得模型需要更新的參數(shù)值
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)


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

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

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

?? ??? ?# 每隔一段時間輸出一個訓(xùn)練結(jié)果
? ? ? ? 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()))

# 訓(xùn)練好的模型預(yù)測測試數(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)