'''  
@Project ：pythonProject @File    ：read_data.py  
@IDE     ：PyCharm @Author  ：周大猛  
@Date    ：2024/05/02 23:50 '''  
from torch.utils.data import Dataset  
from PIL import Image  
import os  
  
class MyDataset(Dataset):  
    def __init__(self, root_dir, label_dir):  
        self.root_dir = root_dir  
        self.label_dir = label_dir  
        self.path = os.path.join(root_dir, label_dir)  
        self.img_path = os.listdir(self.path)  
    def __getitem__(self, index):  
        """  
        读取每一个图片  
        :param index:        :return:  
        """        img_name = self.img_path[index]  
        img_item_path = os.path.join(self.path, img_name)  
        img = Image.open(img_item_path)  
        label = self.label_dir  
        return img, label  
  
    def __len__(self):  
        """获得数据集的长度"""  
        return len(self.img_path)  
  
  
if __name__ == '__main__':  
    root_dir = "dataset/train"  
    ants_label_dir = "ants"  
    bees_label_dir = "bees"  
    ants_dataset = MyDataset(root_dir,ants_label_dir)  
    bees_dataset = MyDataset(root_dir,bees_label_dir)  
  
    # 拼接数据集，按顺序拼接。  
    train_loader = ants_dataset + bees_dataset

class SummaryWriter:  
    """Writes entries directly to event files in the log_dir to be consumed by TensorBoard.  
  
    The `SummaryWriter` class provides a high-level API to create an event file    in a given directory and add summaries and events to it. The class updates the    file contents asynchronously. This allows a training program to call methods    to add data to the file directly from the training loop, without slowing down    training.    """  
    def __init__(  
        self,  
        log_dir=None,  
        comment="",  
        purge_step=None,  
        max_queue=10,  
        flush_secs=120,  
        filename_suffix="",  
    ):  
        """Create a `SummaryWriter` that will write out events and summaries to the event file.  
  
        Args:            log_dir (str): Save directory location. Default is              runs/**CURRENT_DATETIME_HOSTNAME**, which changes after each run.              Use hierarchical folder structure to compare              between runs easily. e.g. pass in 'runs/exp1', 'runs/exp2', etc.              for each new experiment to compare across them.            comment (str): Comment log_dir suffix appended to the default              ``log_dir``. If ``log_dir`` is assigned, this argument has no effect.            purge_step (int):              When logging crashes at step :math:`T+X` and restarts at step :math:`T`,              any events whose global_step larger or equal to :math:`T` will be              purged and hidden from TensorBoard.              Note that crashed and resumed experiments should have the same ``log_dir``.            max_queue (int): Size of the queue for pending events and              summaries before one of the 'add' calls forces a flush to disk.              Default is ten items.            flush_secs (int): How often, in seconds, to flush the              pending events and summaries to disk. Default is every two minutes.            filename_suffix (str): Suffix added to all event filenames in              the log_dir directory. More details on filename construction in              tensorboard.summary.writer.event_file_writer.EventFileWriter.  
        Examples::  
            from torch.utils.tensorboard import SummaryWriter  
            # create a summary writer with automatically generated folder name.            writer = SummaryWriter()            # folder location: runs/May04_22-14-54_s-MacBook-Pro.local/  
            # create a summary writer using the specified folder name.            writer = SummaryWriter("my_experiment")            # folder location: my_experiment  
            # create a summary writer with comment appended.            writer = SummaryWriter(comment="LR_0.1_BATCH_16")            # folder location: runs/May04_22-14-54_s-MacBook-Pro.localLR_0.1_BATCH_16/  
        """

from PIL import Image  
from torch.utils.tensorboard import SummaryWriter  
import numpy as np  
  
writer = SummaryWriter('logs')  
image_path = "dataset/train/ants/5650366_e22b7e1065.jpg"  
iamge_PIL = Image.open(image_path)  
img_array = np.array(iamge_PIL) # 但是这里形状不对,np读出来之后，通道在最后：(375, 500, 3)  
print(img_array.shape)  
  
writer.add_image("image", img_array, 1,dataformats='HWC') # 根据官方文档里面，指定type的顺序

# 查看网络结构的方法  
writer = SummaryWriter("./logs")  
writer.add_graph(net,input)  
writer.close()

from torchvision import transforms

tensor_trans = transforms.ToTensor() # 实例化这个工具  
tensor_img = tensor_trans(image)  # 使用这个工具，输出一个结果

trans_norm = transforms.Normalize([0.5,0.5,0.5], [0.5,0.5,0.5])  
img_norm = trans_norm(img_tensor)  
  
writer.add_image('Normal_img', img_norm)

# Resize  
print(img.size)  
trans_resize = transforms.Resize((512,512))  
# img PIL --> resize --> img_resize PIL  
img_resize = trans_resize(img)  
# img_resize PIL --》to_tensor --> img_resize tensor  
img_resize = trans_totensor(img_resize)  
writer.add_image('Resize_img', img_resize,0)  
print(img_resize)  
writer.close()

# Compose - resize -2  
trans_resize_2 = transforms.Resize(64)  
  
# 这里列表的顺序需要保证前一个的输出类型是后一个的输入类型。  
trans_compose = transforms.Compose([trans_resize_2,trans_totensor,])  
img_resize_2 = trans_compose(img)

# RandomCrop  
trans_random = transforms.RandomCrop(64)  
trans_compose_2 = transforms.Compose([trans_random,trans_totensor,])  
for i in range(10):  
    img_crop = trans_compose_2(img)  
    writer.add_image('Random_img', img_crop,i)

import torchvision  
  
train_set = torchvision.datasets.CIFAR10(root='./data', train=True, download=True)  
test_set = torchvision.datasets.CIFAR10(root='./data', train=False, download=True)

import torchvision  
from torch.utils.tensorboard import SummaryWriter  
  
dataset_transform = torchvision.transforms.Compose([  
  
    torchvision.transforms.ToTensor(),  
    torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),  
  
])  
  
train_set = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=dataset_transform)  
test_set = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=dataset_transform)  
  
print(test_set[0])  
  
writer = SummaryWriter(log_dir='./logs')  
  
for i in range(10):  
    img, label = test_set[i]  
    writer.add_image('test_img', img, i)
    

import torchvision  
from torch.utils.data import DataLoader  
from torch.utils.tensorboard import SummaryWriter  
from torchvision import transforms  
  
data_transforms = transforms.Compose([transforms.ToTensor(),])  
test_data = torchvision.datasets.CIFAR10(root='./data', train=False, download=True,transform=data_transforms)  
  
test_loader = DataLoader(dataset=test_data, batch_size=64, shuffle=True,num_workers=0, drop_last=False)  
  
  
img, label = test_data[0]  
print(img.shape)  
print(label)  
  
writer = SummaryWriter('./logs')  
  
# 多次让dataloader取数据，shuffle就会在每次的epoch影响取值，True会打乱数据集  
for epoch in range(10):  
    step = 0  
    for data in test_loader:  
        imgs, labels = data  
        # 注意这里用的是images  
        writer.add_images('Epoch:{}'.format(epoch), imgs,step)  
        step += 1  
  
writer.close()

import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self):
    # 调用父类的初始化
        super().__init__()
        
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
    # 神经网络的前向传播
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))  # 这个demo进行了两次非线性卷积

import torch  
import torch.nn.functional as F  
  
input = torch.tensor([[1,2,0,3,1],  
                      [0,1,2,3,1],  
                      [1,2,1,0,0],  
                      [5,2,3,1,1],  
                      [2,1,0,1,1]])  
  
kernel = torch.tensor([[1,2,1],  
                       [0,1,0],  
                       [2,1,0]])  
  
# print(input)  
# print(kernel)  
#  
# print(input.shape)  
# print(kernel.shape)  
  
  
# 转换成nn.conv需要的形状  
input = torch.reshape(input,(1,1,5,5)) # shape:batch，通道，长，宽  
kernel = torch.reshape(kernel,(1,1,3,3))  
  
# print(input)  
# print(kernel)  
#  
# print(input.shape)  
# print(kernel.shape)  
  
output1 = F.conv2d(input,kernel,stride=1)  
print(output1)  
  
output2 = F.conv2d(input,kernel,stride=2)  
print(output2)  
  
output3 = F.conv2d(input,kernel,stride=1,padding=1)  
print(output3)

import torch  
import torch.nn as nn  
import torch.nn.functional as F  
import torchvision.datasets as datasets  
from torch.utils.data import DataLoader  
from torch.utils.tensorboard import SummaryWriter  
from torchvision import transforms  
  
dataset = datasets.CIFAR10(root='./data', train=False, download=True,transform=transforms.ToTensor())  
dataloader = DataLoader(dataset, batch_size=64, shuffle=True)  
# input = torch.tensor([  
#     [1,2,0,3,1],  
#     [0,1,2,3,1],  
#     [1,2,1,0,0],  
#     [5,2,3,1,1],  
#     [2,1,0,1,1]  
# ],dtype=torch.float32)  
#  
# input = torch.reshape(input,(-1,1,5,5))  
# print(input.shape)  
  
class T(nn.Module):  
    def __init__(self):  
        super(T, self).__init__()  
        self.pool = nn.MaxPool2d(3,ceil_mode=False)  
  
    def forward(self, x):  
        output = self.pool(x)  
  
        return output  
  
  
tt = T()  
# result = tt(input)  
# print(result.shape)  
# print(result)  
writer = SummaryWriter('./maxpool_logs')  
step = 0  
for data in dataloader:  
    imgs, labels = data  
    writer.add_image('Input', imgs,step,dataformats='NCHW')  
    outputs = tt(imgs)  
    """  
    最大池化不会改变形状，  
    所以不用像卷积那样还要将得到的图片进行reshape  
    """    writer.add_image("Output", outputs,step,dataformats='NCHW')  
  
    step += 1  
  
writer.close()

import torch  
  
i = torch.tensor([[-1, 2, 3], [4, -5, 6], [7, 8, -9]])  
i = torch.reshape(i,(-1,1,3,3))  
  
class DemoModule(torch.nn.Module):  
    def __init__(self):  
        super(DemoModule, self).__init__()  
        self.relu1 = torch.nn.ReLU(inplace=False)  
  
    def forward(self, x):  
        return self.relu1(x)  
  
mod = DemoModule()  
output = mod(i)  
print(output)

import torch.nn as nn  
import torch  
loss = nn.L1Loss()  
input = torch.randn(3, 5, requires_grad=True)  
target = torch.randn(3, 5)  
print(input)  
print(target)  
output = loss(input, target)  
print(output)  
  
loss_mes = nn.MSELoss()  
  
result_mse = loss_mes(input, target)  
print(result_mse)

optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
optimizer = optim.Adam([var1, var2], lr=0.0001)

optim.SGD([
                {'params': model.base.parameters(), 'lr': 1e-2},
                {'params': model.classifier.parameters()}
            ], lr=1e-3, momentum=0.9)

for input, target in dataset:
# - 在进行反向传播和梯度计算之前，先将优化器中的所有参数的梯度缓存清零。
    optimizer.zero_grad()
    output = model(input)
    # 计算模型输出 `output` 和目标标签 `target` 之间的损失（误差）
    loss = loss_fn(output, target)
    # 进行反向传播，计算损失相对于模型参数的梯度
    loss.backward()
    # - 使用计算得到的梯度，按照优化算法更新模型的参数。
    # - 这里的 `optimizer` 是前面定义的优化器（如 `optim.SGD`），它根据参数的梯度和学习率来调整参数的值，使损失函数逐渐减小，从而优化模型。
    optimizer.step()

loss = nn.CrossEntropyLoss()# 交叉熵  
net = NeuralNetwork()  
optimizer = torch.optim.SGD(net.parameters(), lr=0.01)  
for epoch in range(20):  
    running_loss = 0.0 # 每次开始前，把loss设置为0  
    for data in test_loader:  
        """  
        这个for相当于只对data进行了一轮的学习，  
        通常需要好几轮的学习，才能有所改善。        所以需要外层的epoch  
        """        imgs, labels = data  
        outputs = net(imgs)  
  
        result_loss = loss(outputs, labels)  
  
        optimizer.zero_grad() # 对之前的记录清零  
        result_loss.backward()  
        optimizer.step()  
        # print(result_loss)  
        running_loss = running_loss + result_loss  
        # 整体误差总和  
    print("epoch:{}, loss:{}".format(epoch, running_loss))

vgg16_false.add_module('add_linear', nn.Linear(1000, 10))

vgg16_false.classifier.add_module('add_linear', nn.Linear(1000, 10))

vgg16 = torchvision.models.vgg16(pretrained=False)
torch.save(vgg16, 'vgg16_method.pth')

# Load Method 1:  
model_1 = torch.load("vgg16_method.pth")

torch.save(vgg16.state_dict(), 'vgg16_state_dict.pth')

# Specify a path
PATH = "state_dict_model.pt"

# Save
torch.save(net.state_dict(), PATH)

# Load
model = Net()
model.load_state_dict(torch.load(PATH))
model.eval()

'''  
@Project ：pythonProject @File    ：MyModel.py  
@IDE     ：PyCharm @Author  ：周大猛  
@Date    ：2024/05/20 15:21 '''  
import torch  
from torch import nn  
  
  
class TongModel(nn.Module):  
    def __init__(self):  
        super(TongModel, self).__init__()  
        self.model = nn.Sequential(  
            nn.Conv2d(3, 32, kernel_size=5, stride=1, padding=2),  
            nn.MaxPool2d(2),  
            nn.Conv2d(32, 32, 5, 1, 2),  
            nn.MaxPool2d(2),  
            nn.Conv2d(32, 64, 5, 1, 2),  
            nn.MaxPool2d(2),  
            nn.Flatten(),  
            nn.Linear(64 * 4 * 4, 64),  
            nn.Linear(64, 10),  
        )  
  
    def forward(self, x):  
        x = self.model(x)  
        return x  
  
  
if __name__ == '__main__':  
    # 测试网络模型的正确性  
    model = TongModel()  
    # 64个图片，3个通道，尺寸32*32  
    input = torch.ones((64, 3, 32, 32))  
    output = model(input)  
    #torch.Size([64, 10]) 64个图片，10个数据表示每个类别的可能性  
    print(output.shape)

# 准备训练数据  
train_data  = torchvision.datasets.CIFAR10(root='./data', train=True, download=True,transform=transforms.ToTensor())  
  
# 准备测试数据  
test_data = torchvision.datasets.CIFAR10(root='./data', train=False, download=True,transform=transforms.ToTensor())  
  
# 查看数据集信息  
train_data_size = len(train_data)  
test_data_size = len(test_data)  
print("Train data size: {}".format(train_data_size) )  
print("Test data size: %d" % test_data_size)  
  
# 使用dataloader加载数据集  
train_dataloader = DataLoader(train_data, batch_size=64, shuffle=True)  
test_dataloader = DataLoader(test_data, batch_size=64, shuffle=True)

from MyModel import *  
  
# 创建网络模型  
model = TongModel()  
  
# 损失函数和优化器  
loss_fn = nn.CrossEntropyLoss()  
learning_rate = 1e-2  
optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=0.9, weight_decay=5e-4)  
  
# 设置训练参数  
total_train_step = 0 # 记录训练次数  
total_test_step = 0 # 记录测试次数  
epoch = 100 # 训练的轮数  
  
# 使用tensorboard记录数据  
writer = SummaryWriter('./logs_train')

for i in range(epoch):  
    print('Epoch {}/{}'.format(i, epoch))  
  
    # 训练开始  
    model.train() # 当模型有dropout等特殊层的时候，起作用  
    for data in train_dataloader:  
        imgs, labels = data  
        output = model(imgs)  
        # 计算输出和真实的损失  
        loss = loss_fn(output, labels)  
  
        # 优化器优化模型  
        optimizer.zero_grad() # 准备优化，先梯度清零  
        loss.backward() #得到每个参数的梯度  
        optimizer.step() # 对参数进行优化  
        total_train_step += 1  
        if total_train_step % 100 == 0:  
            print('Total train step:{}, Loss:{}' .format(total_train_step, loss.item()))  
            writer.add_scalar('Loss/train', loss.item(), total_train_step)  
    """  
    每次训练一轮之后，需要知道本次训练之后在测试集上模型表现是否有进步。  
    在第一层的for中，对此进行检测。  
    这里不需要对模型进行调优。    需要知道在整个数据集上的loss  
    """    model.eval() # 与train（）一个情况  
    total_test_loss = 0.0  
  
    # 计算ACC  
    total_accuracy = 0  
    with torch.no_grad():  
        # 没有梯度了  
        # 测试开始  
        for data in test_dataloader:  
            imgs, labels = data  
            output = model(imgs)  
            loss = loss_fn(output, labels)  
            total_test_step += 1  
            total_test_loss += loss.item()  
            accuracy = (output.argmax(dim=1)==labels).sum()  
            total_accuracy += accuracy.item()  
  
        print('Total test  Loss:{}' .format(total_test_loss))  
        writer.add_scalar('Loss/test', total_test_loss, total_test_step)  
        print("整体测试集的正确率:{}".format(total_accuracy/test_data_size))  
        writer.add_scalar('Accuracy/test', total_accuracy/test_data_size, total_test_step)  
  
    torch.save(model.state_dict(), './modelData/model_{}.pt'.format(i))  
writer.close()

# 创建网络模型  
model = TongModel()  
model = model.cuda()

imgs, labels = data  
imgs = imgs.cuda()  
labels = labels.cuda()

# 损失函数和优化器  
loss_fn = nn.CrossEntropyLoss()  
loss_fn = loss_fn.cuda()

model.to(device)