import torch
from torchvision import datasets
from torchvision.transforms import ToTensor
import matplotlib.pyplot as plt
import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'	# 沒有這句會報錯，具體原因我也不知道

training_data = datasets.FashionMNIST(
    root="../data",
    train=True,
    download=True,
    transform=ToTensor()
)

test_data = datasets.FashionMNIST(
    root="../data",
    train=False,
    download=True,
    transform=ToTensor()
)

labels_map = {
    0: "T-Shirt",
    1: "Trouser",
    2: "Pullover",
    3: "Dress",
    4: "Coat",
    5: "Sandal",
    6: "Shirt",
    7: "Sneaker",
    8: "Bag",
    9: "Ankle Boot",
}
figure = plt.figure(figsize=(8, 8))
cols, rows = 3, 3
for i in range(1, cols * rows + 1):
    # 從0-len(training_data)中隨機(jī)生成一個數(shù)字（不包括右邊界）
    sample_idx = torch.randint(len(training_data), size=(1,)).item() 
    img, label = training_data[sample_idx] # 獲得圖片和標(biāo)簽
    figure.add_subplot(rows, cols, i)
    plt.title(labels_map[label])
    plt.axis("off")                        # 坐標(biāo)軸不可見
    plt.imshow(img.squeeze(), cmap="gray") # 顯示灰度圖
plt.show()

import os
import pandas as pd
from torch.utils.data import Dataset
from torchvision.io import read_image

class CustomImageDataset(Dataset):
    def __init__(self, annotations_file, img_dir, transform=None, target_transform=None):
        self.img_labels = pd.read_csv(annotations_file)
        self.img_dir = img_dir
        self.transform = transform
        self.target_transform = target_transform

    def __len__(self):
        return len(self.img_labels)

    def __getitem__(self, idx):
        img_path = os.path.join(self.img_dir, self.img_labels.iloc[idx, 0])
        image = read_image(img_path)
        label = self.img_labels.iloc[idx, 1]
        if self.transform:
            image = self.transform(image)
        if self.target_transform:
            label = self.target_transform(label)
        return image, label

from torch.utils.data import DataLoader

train_dataloader = DataLoader(training_data, batch_size=64, shuffle=True) 
# train_dataloader是一個DataLoader類的對象
test_dataloader = DataLoader(test_data, batch_size=64, shuffle=True)

for batch_index, (features, label) in enumerate(train_dataloader):
    print(batch_index)
    print(f"Feature batch shape: {features.size()}")
    print(f"Labels batch shape: {label.size()}")
    img = features[0].squeeze()
    label = label[0]
    plt.imshow(img, cmap="gray")
    plt.show()
    print(f"Label: {label}")
    break

for batch_index, (features, label) in enumerate(train_dataloader):
    print(batch_index)
    print(f"Feature batch shape: {features.size()}")
    print(f"Labels batch shape: {label.size()}")
    img = features[0].squeeze()
    label = label[0]
    plt.imshow(img, cmap="gray")
    plt.show()
    print(f"Label: {label}")
    break

import torch
from torchvision import datasets
from torchvision.transforms import ToTensor, Lambda

ds = datasets.FashionMNIST(
    root="data",
    train=True,
    download=True,
    transform=ToTensor(),
    target_transform=Lambda(lambda y: torch.zeros(10, dtype=torch.float).scatter_(dim=0, index=torch.tensor(y), value=1))
)

class NeuralNetwork(nn.Module):
    def __init__(self):
        super(NeuralNetwork, self).__init__()
        self.flatten = nn.Flatten()				# 打平
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28*28, 512),		# 線性層
            nn.ReLU(),					# 激活層
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 10),
            nn.ReLU()
        )

    def forward(self, x):
        x = self.flatten(x)					# 打平層
        logits = self.linear_relu_stack(x)	# 線性激活層
        return logits

model = NeuralNetwork()
print(model)

X = torch.rand(1, 28, 28)	# 生成一個樣本
logits = model(X)							# 將樣本輸入到模型中，自動調(diào)用forward方法
print(logits.size())						
pred_probab = nn.Softmax(dim=1)(logits)		# 實(shí)例化一個Softmax對象，并通過對象調(diào)用
y_pred = pred_probab.argmax(1)				# 獲得概率最大索引
print(f"Predicted class: {y_pred}")

input_image = torch.rand(3,28,28)
print(input_image.size())

layer1 = nn.Linear(in_features=28*28, out_features=20)
hidden1 = layer1(flat_image)
print(hidden1.size())

w = torch.rand(784, 20)
b = torch.rand((1, 20))
hidden2 = flat_image @ w + b
print(hidden2.size())

seq_modules = nn.Sequential(
    flatten,		# 在nn.Sequential中可以調(diào)用其他模塊
    layer1,
    nn.ReLU(),
    nn.Linear(20, 10)
)
input_image = torch.rand(3,28,28)
logits = seq_modules(input_image)

para_list = list(model.parameters())

# 將參數(shù)生成器轉(zhuǎn)換成列表之后，列表的第一個元素是w，第二個元素是b
print(type(para_list[0]))
print(f'number of linear_layers :{len(para_list)/2}')
print('weights:')
print(para_list[0][:2])  # 只切出兩個樣本來顯示
print('bias:')
print(para_list[1][:2]) 
print('\nthe shape of first linear layer:', para_list[0].shape)

print("Model structure: ", model, "\n\n")

for name, param in model.named_parameters():
    print(f"Layer: {name} | Size: {param.size()} | Values : \n{param[:2]} \n")
    # 只切出前兩行顯示

learning_rate = 1e-3
batch_size = 64
epochs = 5

# Initialize the loss function
loss_fn = nn.CrossEntropyLoss()

optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

def train_loop(dataloader, model, loss_fn, optimizer):
    size = len(dataloader.dataset)
    for batch, (X, y) in enumerate(dataloader):
        # Compute prediction and loss，下面兩步相當(dāng)于繪制計(jì)算圖
        pred = model(X)
        loss = loss_fn(pred, y)

        # Backpropagation
        optimizer.zero_grad()	# 梯度信息清零
        loss.backward()			# 反向傳播
        optimizer.step()		# 一旦有了梯度，就可以更新參數(shù)

        if batch % 100 == 0:
            loss, current = loss.item(), batch * len(X)
            print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")


def test_loop(dataloader, model, loss_fn):
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    test_loss, correct = 0, 0

    with torch.no_grad():		# 禁用梯度跟蹤，后面會講
        for X, y in dataloader:
            pred = model(X)
            test_loss += loss_fn(pred, y).item()
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()

    test_loss /= num_batches
    correct /= size
    print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")

for t in range(epochs):
    print(f"Epoch {t+1}\n-------------------------------")
    train_loop(train_dataloader, model, loss_fn, optimizer)
    test_loop(test_dataloader, model, loss_fn)
print("Done!")

test_features, test_labels = next(iter(test_dataloader))
logits = model(test_features[0])
pred_probab = nn.Softmax(dim=1)(logits)
pred = pred_probab.argmax(1)

img = test_features[0].squeeze()    # 將長度為1的維度去掉
true_label = test_labels[0]
plt.imshow(img, cmap="gray")
plt.show()
print(f"True label: {true_label}")
print(f"Predict: {pred.item()}")

z = torch.matmul(x, w)+b
print(z.requires_grad)

with torch.no_grad():
    z = torch.matmul(x, w)+b
	print(z.requires_grad)

# 保存模型參數(shù)
torch.save(model.state_dict(), './model_weights.pth')
# 保存模型結(jié)構(gòu)
torch.save(model, './model.pth')

# 導(dǎo)入模型結(jié)構(gòu)
model_loaded = torch.load('./model.pth') 
# 如果原來的model在cuda:0上，那么導(dǎo)入之后，model_loaded也在cuda:0上

# 導(dǎo)入模型參數(shù)
model_loaded.load_state_dict(torch.load('model_weights.pth'))

# 實(shí)例化一個NeuralNetwork對象
model_loaded = NeuralNetwork() # model_loaded默認(rèn)是CPU中
# 導(dǎo)入模型參數(shù)
model_loaded.load_state_dict(torch.load('./gdrive/MyDrive/model_weights.pth'))

test_features, test_labels = next(iter(test_dataloader))
logits = model_loaded(test_features[0])   # 使用導(dǎo)入的模型
pred_probab = nn.Softmax(dim=1)(logits)
pred = pred_probab.argmax(1)

img = test_features[0].squeeze()    # 將長度為1的維度去掉
true_label = test_labels[0]
plt.imshow(img, cmap="gray")
plt.show()
print(f"True label: {true_label}")
print(f"Predict: {pred.item()}")

# We can move our tensor to the GPU if available
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print('Using {} device'.format(device))

model = NeuralNetwork().to(device)

X = torch.rand(1, 28, 28, device=device)	# 創(chuàng)建時指定設(shè)備
Y = torch.rand(10).to(device)				# 創(chuàng)建后轉(zhuǎn)移

device = torch.device(‘cuda:2')
# 2是設(shè)備號，假如有八張顯卡，那么編號就是0—7

# coding=utf-8
import torch
import torch.nn as nn
from torchvision import datasets
from torch.utils.data import DataLoader
from torchvision.transforms import ToTensor
import matplotlib.pyplot as plt
import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'	# 沒有這句會報錯，具體原因我也不知道

# 導(dǎo)入數(shù)據(jù)
training_data = datasets.FashionMNIST(
    root="../data",
    train=True,
    download=True,
    transform=ToTensor()
)
test_data = datasets.FashionMNIST(
    root="../data",
    train=False,
    download=True,
    transform=ToTensor()
)

# 定義超參數(shù)
# 之所以在這個地方定義，是因?yàn)樵诔跏蓟疍ataLoader時需要用到batch_size
learning_rate = 1e-3
batch_size = 64
epochs = 5
train_dataloader = DataLoader(training_data, batch_size=batch_size, shuffle=True)
# train_dataloader是一個DataLoader類的對象
test_dataloader = DataLoader(test_data, batch_size=batch_size, shuffle=True)


# 搭建神經(jīng)網(wǎng)絡(luò)
class NeuralNetwork(nn.Module):
    def __init__(self):
        super(NeuralNetwork, self).__init__()
        self.flatten = nn.Flatten()				# 打平
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28*28, 512),		# 線性層
            nn.ReLU(),					# 激活層
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 10),
            nn.ReLU()
        )

    def forward(self, x):
        x = self.flatten(x)					# 打平層
        logits = self.linear_relu_stack(x)	# 線性激活層
        return logits

# 確定使用設(shè)備
device = 'cuda' if torch.cuda.is_available() else 'cpu'

# 實(shí)例化一個神經(jīng)網(wǎng)絡(luò)類
model = NeuralNetwork().to(device)

# 確定損失函數(shù)
loss_fn = nn.CrossEntropyLoss()

# 確定優(yōu)化器
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)


# 封裝訓(xùn)練過程
def train_loop(dataloader, model, loss_fn, optimizer):
    size = len(dataloader.dataset)
    for batch, (X, y) in enumerate(dataloader):
        X, y = X.to(device), y.to(device) # 將樣本和標(biāo)簽轉(zhuǎn)移到device中
        # Compute prediction and loss，下面兩步相當(dāng)于繪制計(jì)算圖
        pred = model(X)
        loss = loss_fn(pred, y)

        # Backpropagation
        optimizer.zero_grad()	# 梯度信息清零
        loss.backward()			# 反向傳播
        optimizer.step()		# 一旦有了梯度，就可以更新參數(shù)

        if batch % 100 == 0:
            loss, current = loss.item(), batch * len(X)
            print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")
            

# 封裝測試過程
def test_loop(dataloader, model, loss_fn):
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    test_loss, correct = 0, 0

    with torch.no_grad():
        for X, y in dataloader:
            X, y = X.to(device), y.to(device)  # 將樣本和標(biāo)簽轉(zhuǎn)移到device中
            pred = model(X)
            test_loss += loss_fn(pred, y).item()
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()

    test_loss /= num_batches
    correct /= size
    print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")

# 訓(xùn)練
for t in range(epochs):
    print(f"Epoch {t+1}\n-------------------------------")
    train_loop(train_dataloader, model, loss_fn, optimizer)
    test_loop(test_dataloader, model, loss_fn)
print("Done!")

# 保存模型參數(shù)
torch.save(model.state_dict(), './model_weights.pth')
# 保存模型結(jié)構(gòu)
torch.save(model, './model.pth')

# 導(dǎo)入模型結(jié)構(gòu)
model_loaded = torch.load('./model.pth')    # 模型自動導(dǎo)入到GPU當(dāng)中
# 導(dǎo)入模型參數(shù)
model_loaded.load_state_dict(torch.load('model_weights.pth'))

# 用導(dǎo)入的模型測試
test_features, test_labels = next(iter(test_dataloader))
test_features = test_features.to(device)
logits = model_loaded(test_features[0])   # 使用導(dǎo)入的模型
pred_probab = nn.Softmax(dim=1)(logits)
pred = pred_probab.argmax(1)

# 可視化
img = test_features[0].squeeze()    # 將長度為1的維度去掉
img = img.to('cpu')              # 繪圖時，需要將張量轉(zhuǎn)回到CPU當(dāng)中
true_label = test_labels[0]   # 標(biāo)簽是否轉(zhuǎn)移到CPU無所謂，因?yàn)闆]有用于plt的方法中
plt.imshow(img, cmap="gray")
plt.show()
print(f"True label: {true_label}")
print(f"Predict: {pred.item()}")