1. 图像增广

# -*- coding: utf-8 -*-
# @Project: zc
# @Author: ZhangChu
# @File name: torchvision_net
# @Create time: 2021/12/14 10:28

# 0. 导入相关数据库
import matplotlib.pyplot as plt
import torch
import torchvision
from d2l import torch as d2l

# 1. 从datasets里面加载 CIFAR10 的数据
all_imags = torchvision.datasets.CIFAR10(train=True, root="../data_new", download=True)
# print(all_imags)
# d2l.show_images([all_imags[i][0] for i in range(32)],4,8,scale=0.8)

# 2. 定义 (训练) 数据集的图像增广的方式，用 Compose进行组合
train_augs = torchvision.transforms.Compose([
	torchvision.transforms.RandomHorizontalFlip(),
	torchvision.transforms.ToTensor()
])

# 3. 定义 (测试) 数据集图像处理方式，一般测试集不会做图像增广，只是转换成张量即可
test_augs = torchvision.transforms.Compose([
	torchvision.transforms.ToTensor()
])


# 4. 加载测试集和训练集，转换成 DataLoader
def load_cifar10(is_train, augs, batch_size):
	dataset = torchvision.datasets.CIFAR10(root="../data", train=is_train,
										   transform=augs, download=True)
	dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size,
											 shuffle=is_train, num_workers=d2l.get_dataloader_workers())
	return dataloader


# 5. 用多 GPU 进行小批量训练
def train_batch_ch13(net, X, y, loss, trainer, devices):
	"""
	 function :用多 GPU 进行小批量训练
	:param net: 神经网络 net
	:param X: 输入样本 X
	:param y: 实际标签 Y
	:param loss: 损失函数
	:param trainer: 迭代器
	:param devices: GPU 或 CPU
	:return: train_loss_sum 训练损失值，train_acc_sum 训练精度值
	"""
	if isinstance(X, list):  # 判断 X 是否是列表 list
		X = [x.to(devices[0]) for x in X]  # 将 X 迭代放到 GPU上
	else:
		X = X.to(devices[0])
	y = y.to(devices[0])  # 将 Y 放到 GPU 上
	net.train()  # 神经网络训练模式
	trainer.zero_grad()  # 迭代器清零
	pred = net(X)  # 得到通过 神经网络 NET 后的 预测值 pred
	l = loss(pred, y)  # 将预测 pred 和真实标签 y 进行对比，计算损失
	l.sum().backward()  # 损失求和得到标量后进行回传
	trainer.step()  # 迭代器更新数值
	train_loss_sum = l.sum()  # 将损失值赋值给 train_loss_sum ，方便后续回传
	train_acc_sum = d2l.accuracy(pred, y)  # 计算精度
	return train_loss_sum, train_acc_sum  # 返回训练损失和训练精度
plt.show()