import numpy as np
import torch

def step_func(x): # 参数x可以接受numpy数组
    y = x >= 0 # y是布尔型数组
    return y.type(torch.int)
# 定义单个神经元 2输入
net = torch.nn.Sequential(torch.nn.Linear(2,1)) 

# 定义权重
w = torch.tensor([1,1]).type(torch.float32)
b = -2

net[0].weight.data = w.reshape(net[0].weight.shape)
net[0].bias.data.fill_(b)

# 验证输出
x = torch.tensor([1.0, 1.0]).type(torch.float32)
print('The result of 1 and 1 is', step_func(net.forward(x)).item())
x = torch.tensor([0, 1.0]).type(torch.float32)
print('The result of 1 and 1 is', step_func(net.forward(x)).item())
x = torch.tensor([0, 0]).type(torch.float32)
print('The result of 1 and 1 is', step_func(net.forward(x)).item())

torch.matmul(x, w) + b

from torchvision import transforms # 图像变换、数据预处理
import cv2
import matplotlib.pyplot as plt
#Compose是一个容器，传入的参数是列表，ToTensor()，类型变换，Normalize是数据标准化，去均值，除标准差
# transforms.ToTensor() 把取值[0, 255]的PIL图像形状为[H,W,C]转换为形状[C,H,W] 取值范围[0,1.0]的张量
# transforms.RandomRotation(degrees=(0, 180))
transform = transforms.Compose([transforms.ToTensor(),transforms.RandomPerspective(distortion_scale=0.6, p=1.0),transforms.RandomRotation(degrees=(0, 180))])

cat_bgr = cv2.imread(r'./https://haodong-li.com/zju/assets/img/vision_2022/cat.jpg')
cat_rgb = cv2.cvtColor(cat_bgr, cv2.COLOR_BGR2RGB)
plt.subplot(1,2,1)
plt.imshow(cat_rgb)
cat_trans = transform(cat_rgb)
plt.subplot(1,2,2)
plt.imshow(cat_trans.numpy().transpose((1,2,0)))

import numpy as np
import torch
# 定义单个神经元 2输入
net = torch.nn.Sequential(torch.nn.Linear(2,2),torch.nn.Linear(2,1)) 

# 定义权重
net[0].weight.data = torch.tensor([1,2,3,4]).type(torch.float32).reshape(net[0].weight.shape)
net[1].weight.data = torch.tensor([0.5,0.6]).type(torch.float32).reshape(net[1].weight.shape)

net[0].bias.data.fill_(0)
net[1].bias.data.fill_(0)

x = torch.tensor([1.0, 0.5]).type(torch.float32)
net(x)

# 数值求导举例
# 定义权重
net[0].weight.data = torch.tensor([0.5,1.5,2.3,3]).type(torch.float32).reshape(net[0].weight.shape)
net[1].weight.data = torch.tensor([1,1]).type(torch.float32).reshape(net[1].weight.shape)
# y1 = 0.5*(4-net(x).item())**2
y1 = net(x).item()
dx = 0.001
net[0].weight.data = torch.tensor([0.5,1.5,2.3,3]).type(torch.float32).reshape(net[0].weight.shape)
net[1].weight.data = torch.tensor([1+dx,1]).type(torch.float32).reshape(net[1].weight.shape)
# y2 = 0.5*(4-net(x).item())**2
y2 = net(x).item()
(y2-y1)/dx

# 数值求导举例 y = x^T * x
import torch
x = torch.arange(4.0)
x.requires_grad_(True)
y = 2 * torch.dot(x, x)

y.backward()
x.grad,x.grad == 4 * x

import torch
def softmax(X):
    X_exp = torch.exp(X)
    partition = X_exp.sum(1, keepdim=True)
    return X_exp / partition  # 这里应用了广播机制
X = torch.normal(0, 1, (2, 5))
X_prob = softmax(X)
def cross_entropy(y_hat, y):
    return - torch.log(y_hat[range(len(y_hat)), y])

y = torch.tensor([0, 2])
y_hat = torch.tensor([[0.1, 0.3, 0.6], [0.3, 0.2, 0.5]])


X_prob, X_prob.sum(1),y_hat[[0, 1], y],cross_entropy(y_hat, y)

!pip install torchsummary

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import time
import torch.optim as optim

import torchvision
import torch.utils.data as Data
import torchvision.transforms as transforms
from torchvision.datasets import MNIST,FashionMNIST
print('GPU useable state: ',torch.cuda.is_available())

transform = transforms.Compose(
    [transforms.ToTensor()]) 
#      transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])



# MNIST
data_train = MNIST(root = "data/", ## 数据的路径，如果存在数据则加载，否则下载至此路径
                            transform = transform, ## 图像变换操作
                            train =True, ## 决定使用训练集还是测试集
                            download = True) ## 选择是否需要下载数据
data_test = MNIST(root = "data/",
                            transform = transform,
                            train =False)
## 定义数据加载器
batch_size = 256
train_loader = Data.DataLoader(dataset = data_train, ## 使用的数据集
                                            batch_size = batch_size, ## 批处理样本大小
                                            shuffle = True ## 是否打乱数据顺序
                                            )
test_loader = Data.DataLoader(dataset = data_test,
                                            batch_size = batch_size,
                                            shuffle = True  #将顺序随机打乱
                                            )

class ConvNet(nn.Module):
    def __init__(self, ch = 3, h = 32, w = 32):
        super(ConvNet, self).__init__()
        # 3 input image channel, 6 https://haodong-li.com/zju/assets/img/vision_2022/6/output channels, 5x5 square convolution
        # kernel
        self.conv1 = nn.Sequential(
            nn.Conv2d(
                in_channels = ch, 
                out_channels = 6, 
                kernel_size = 5), ## (N*1* h * w ) ==> (N*6* (h-4) * (w-4))
            nn.ReLU(),            
            nn.AvgPool2d(
                kernel_size = 2,
                stride = 2)       ## (N*6* (h-4) * (w-4)) ==> (N*6* (w-4)/2 * (w-4)/2)
        )
        
        self.conv2 = nn.Sequential(
            nn.Conv2d(6, 16, 5),  ## (N*6* (h-4)/2 * (w-4)/2) ==> (N*16* (h-4)/2-4 * (w-4)/2-4)
            nn.ReLU(),
            nn.AvgPool2d(2,2)     ## (N*16* (h-4)/2-4 * (w-4)/2-4) ==> (N*16* h/4 -3 * w/4 -3)
        )
        
        self.classifier = nn.Sequential(
           nn.Linear(int(16*(h/4-3)*(w/4-3)),120), # 5
           nn.ReLU(),
           nn.Linear(120,84),
           nn.ReLU(),
           nn.Linear(84,10)) 

    def forward(self, x):
        # Max pooling over a (2, 2) window
        x = self.conv1(x)
        # If the size is a square, you can specify with a single number
        x = self.conv2(x)
        x = torch.flatten(x, 1) # flatten all dimensions except the batch dimension
        x = self.classifier(x)
        return x

class dnn(nn.Module):
    def __init__(self):
        super(dnn,self).__init__()
        ## 定义层
        self.flatten = nn.Flatten()
        self.hiden1 = nn.Linear(in_features = 784, out_features = 512, bias = True)
        self.active1 = nn.ReLU() ## 激活函数
        self.drop1 = nn.Dropout(p=0.2)
        self.hiden2 = nn.Linear(512,512)
        self.active2 = nn.ReLU()
        self.drop2 = nn.Dropout(p=0.2)
        self.fc = nn.Linear(512,10)
    def forward(self,x):
        x = self.flatten(x)
        x = self.hiden1(x)
        x = self.active1(x)
        x = self.drop1(x)
        x = self.hiden2(x)
        x = self.active2(x)
        x = self.drop2(x)
        x = self.fc(x)
        return x

net = dnn()
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
net.to(device)
print(device)
from torchsummary import summary
summary(net,(1,28,28))

net = ConvNet(1,28,28)
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
net.to(device)
print(device)
from torchsummary import summary
summary(net,(1,28,28))

## 初始化模型参数
def init_weights(m):
    if type(m) == nn.Conv2d:
        nn.init.normal_(m.weight, mean=0, std=0.5)
    if type(m) == nn.Linear:
        nn.init.normal_(m.weight, std=0.01)

net.apply(init_weights);

import pandas as pd

def train_model(model, train_loader, test_loader, loss_func, optimizer, device, num_epochs = 5):    
    """
    model: 网络模型；     train_loader: 训练数据集； test_loader: 测试数据集
    loss_func: 损失函数； optimizer: 优化方法；      num_epochs: 训练的轮数
    device: 控制是否使用GPU
    """
    train_loss_all = []
    train_acc_all = []
    val_acc_all = []
    val_loss_all = []
    
    gap = len(train_loader)//30
    start_time = time.time()
    
    for epoch in range(num_epochs):
        if epoch:
            print('-'*10)
        print("Epoch {}/{}".format(epoch + 1,num_epochs))
        
        ## 
        train_loss = 0.0
        train_corrects = 0
        train_num = 0
        
        val_loss = 0.0
        val_corrects = 0
        val_num = 0
        
        for step,data in enumerate(train_loader):
            model.train() 
            x,y = data[0].to(device), data[1].to(device)    
            output = model(x) ## 模型在 X 上的输出: N * num_class
            pre_lab = torch.argmax(output, 1) ## 获得预测结果
            loss = loss_func(output, y) ## 损失
            optimizer.zero_grad() ## 每次迭代将梯度初始化为0
            loss.backward() ## 损失的后向传播， 计算梯度
            optimizer.step() ## 使用梯度进行优化
            train_loss += loss.item() * x.size(0) ## 统计模型预测损失
            train_corrects += torch.sum(pre_lab == y.data)
            train_num += x.size(0)
            
            if step %gap  == gap-1:
                cont = step//gap
                if cont >30:
                    cont = 30
                print(str(cont)+'/','30','['+'='*cont+'>'+'-'*(30-cont)+']',
                      'loss: {:.4f} - accuracy: {:.4f}'.format(train_loss/train_num,train_corrects.double().item()/train_num)
                      ,end="\r")
 
        for data in test_loader:
            model.eval()
            X_test, y_test = data[0].to(device), data[1].to(device)     
            X_test.requires_grad=True
            with torch.no_grad():
                output = model(X_test)
            test_loss = loss_func(output, y_test )
            _, pred = torch.max(output.data, 1)
            val_corrects += torch.sum(pred == y_test.data)
            val_loss += test_loss.item()*X_test.size(0)
            val_num += X_test.size(0)
   
        ##
        train_loss_all.append(train_loss/train_num)
        train_acc_all.append(train_corrects.double().item()/train_num)
        val_loss_all.append(val_loss/val_num)
        val_acc_all.append(val_corrects.double().item()/val_num)
        
        print('')
        print("No.{} Train Loss is:{:.4f}, Train_accuracy is {:.4f}%"
              .format(epoch+1, train_loss_all[-1],train_acc_all[-1] * 100))
        print("No.{} Val Loss is:{:.4f},  Val_accuracy is {:.4f}%"
              .format(epoch+1, val_loss_all[-1], val_acc_all[-1] * 100))    
        
        time_use = time.time() - start_time
        print("Train and val complete in {:.0f}m {:.0f}s".format(time_use//60, time_use%60))
        
    train_process = pd.DataFrame(
        data = {"epoch":range(num_epochs),
                "train_loss":train_loss_all,
                "train_acc":train_acc_all,
                "val_loss":val_loss_all,
                "val_acc":val_acc_all})
    return model, train_process

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.1, momentum=0.9)

net, train_process = train_model(net, train_loader, test_loader, criterion, optimizer, device, num_epochs = 10)

# plotting the metrics
fig = plt.figure(figsize=(10,5))
plt.subplot(2,1,1)
plt.plot(train_process.train_acc, "b-")
plt.plot(train_process.val_acc, "r-")
plt.title('model accuracy')
plt.ylabel('accuracy')
plt.xlabel('epoch')
plt.legend(['train', 'test'], loc='lower right')

plt.subplot(2,1,2)
plt.plot(train_process.train_loss, "b-")
plt.plot(train_process.val_loss, "r-")
plt.title('model loss')
plt.ylabel('loss')
plt.xlabel('epoch')
# plt.legend(['train'], loc='upper right')
plt.legend(['train', 'test'], loc='upper right')

plt.tight_layout()

for x,y in test_loader:
    break
plt.imshow(x[0].squeeze().numpy(),'gray')

output = net.conv1(x[0].unsqueeze(dim=0).to(device))
for i in range(6):
    plt.subplot(2,3,i+1)
    plt.imshow(output[0][i].squeeze().detach().cpu().numpy(),'gray')