PyTorch入门

倞涼諒

已于 2022-06-08 09:37:13 修改

阅读量147

点赞数

分类专栏：人工智能文章标签： pytorch 深度学习 python

于 2022-06-08 09:09:05 首次发布

本文链接：https://blog.csdn.net/weixin_44242764/article/details/125177087

版权

人工智能专栏收录该内容

3 篇文章 0 订阅

订阅专栏

视频教程推荐（主页有数据集链接）：PyTorch深度学习快速入门教程（绝对通俗易懂！）【小土堆】
入门学习必备：PyTorch官网
以下代码根据小土堆视频跟学所写

1、准备数据集

本地数据集
构建自己的Dataset类要继承Dataset，且必须重载__getitem__()函数
__getitem__()：通过指定键获取数据

import os
from torch.utils.data import Dataset
from PIL import Image


class MyData(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, idx):
        img_name = self.img_path[idx]
        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)


root_dir = "dataset/train"
bees_label_dir = "bees"
bees_data = MyData(root_dir, bees_label_dir)

在线数据集
PyTorch框架在torchvision.datasets中提供了许多图像数据集，例如

import torchvision

test_set = torchvision.datasets.CIFAR10(root="dataset_transforms", train=False, transform=transforms.ToTensor(), download=True)

2、加载数据集

参数：

dataset(Dataset)：要加载的数据集
batch_size(int)：每批数据加载多少（如从牌堆中每次抽多少张牌），默认为1
shuffle(bool)：在每次epoch是否打乱（如每次发牌是否重新洗牌），默认为False
drop_last(bool)：是否略去最后不能组成batch_size的数据（如54 % 8 = 6，如果drop_last=True，最后6个数据将忽略，否则这6个数据成为一个batch）

data_loader = DataLoader(dataset=test_set, batch_size=64, shuffle=True, drop_last=True)

3、模型

import torch
from torch import nn
from torch.nn import Linear, Flatten, MaxPool2d, Conv2d, Sequential

class Module(nn.Module):
    def __init__(self):
        super(Module, self).__init__()
        self.model = Sequential(
            Conv2d(3, 32, 5, padding=2),
            MaxPool2d(2),
            Conv2d(32, 32, 5, padding=2),
            MaxPool2d(2),
            Conv2d(32, 64, 5, padding=2),
            MaxPool2d(2),
            Flatten(),
            Linear(1024, 64),
            Linear(64, 10)
        )

    def forward(self, input):
        output = self.model(input)
        return output

if __name__ == "__main__":
    module = Module()
    input = torch.ones((64, 3, 32, 32), dtype=torch.float32)
    output = module(input)
    print(output.shape)

4、损失函数

import torch
from torch import nn

inputs = torch.tensor([1, 2, 3], dtype=torch.float32)
targets = torch.tensor([1, 2, 5], dtype=torch.float32)


# MAE
loss = nn.L1Loss()
result = loss(inputs, targets)
print(result)

# MSE
loss_mse = nn.MSELoss()
result_mse = loss_mse(inputs, targets)
print(result_mse)

5、优化器

以SGD为例

module = Module()
optim = torch.optim.SGD(module.parameters(), lr=0.01)

6、模型训练

import time

import torch.optim
import torchvision
from torch import nn
from torch.nn import CrossEntropyLoss
from torch.utils.tensorboard import SummaryWriter
from torch.nn import Linear, Flatten, MaxPool2d, Conv2d, Sequential
from torch.utils.data import DataLoader


# 定义训练设备
device = torch.device("cuda")

# 准备数据集
train_dataset = torchvision.datasets.CIFAR10(root="dataset_transforms", train=True, transform=torchvision.transforms.ToTensor(),
                                             download=True)
test_dataset = torchvision.datasets.CIFAR10(root="dataset_transforms", train=False, transform=torchvision.transforms.ToTensor(),
                                             download=True)

test_dataset_size = len(test_dataset)
# 加载数据集
train_loader = DataLoader(train_dataset, batch_size=64)
test_loader = DataLoader(test_dataset, batch_size=64)

# 搭建神经网络
class Module(nn.Module):
    def __init__(self):
        super(Module, self).__init__()
        self.model = Sequential(
            Conv2d(3, 32, 5, padding=2),
            MaxPool2d(2),
            Conv2d(32, 32, 5, padding=2),
            MaxPool2d(2),
            Conv2d(32, 64, 5, padding=2),
            MaxPool2d(2),
            Flatten(),
            Linear(1024, 64),
            Linear(64, 10)
        )

    def forward(self, input):
        output = self.model(input)
        return output
module = Module()
if torch.cuda.is_available():
    # 方法一
    # module = module.cuda()
    # 方法二
    module = module.to(device)

# 损失函数
loss_func = CrossEntropyLoss()
if torch.cuda.is_available():
    loss_func = loss_func.cuda()

# 优化器
learing_rate = 1e-2
optim = torch.optim.SGD(module.parameters(), lr=learing_rate)

# 记录训练次数
total_train_step = 0
# 记录测试次数
total_test_step = 0
# 训练轮数
epoch = 10


writer = SummaryWriter("logs_train_gpu")

for i in range(epoch):
    print("----------第{}轮训练----------".format(i + 1))
    start_time = time.time()
    # 模型训练步骤
    total_train_loss = 0
    for data in train_loader:
        imgs, targets = data
        if torch.cuda.is_available():
            imgs = imgs.cuda()
            targets = targets.cuda()
        outputs = module(imgs)
        loss = loss_func(outputs, targets)
        # 更新参数
        optim.zero_grad()
        loss.backward()
        optim.step()
        # 统计该轮训练的loss
        total_train_loss += loss
        # 训练次数
        total_train_step += 1

        if total_train_step % 100 == 0:
            print("第{}次的loss为 {}".format(total_train_step, loss.item()))
            end_time = time.time()
            print(end_time - start_time)
            writer.add_scalar("train_loss", loss.item(), total_train_step)

    # 模型测试步骤
    total_test_loss = 0
    total__accuracy = 0
    with torch.no_grad():
        for data in test_loader:
            imgs, targets = data
            if torch.cuda.is_available():
                imgs = imgs.cuda()
                targets = targets.cuda()
            outputs = module(imgs)
            loss = loss_func(outputs, targets)
            total_test_loss += loss
            accuracy = (outputs.argmax(1) == targets).sum()
            total__accuracy += accuracy

    print("测试集上的loss为 {}".format(total_test_loss))
    print("测试集上的准确率为 {}".format(total__accuracy / test_dataset_size))
    writer.add_scalar("test_loss", total_test_loss, total_test_step)
    writer.add_scalar("accuracy_rate", total__accuracy / test_dataset_size, total_test_step)
    total_test_step += 1

    torch.save(module, "module{}.pth".format(i))
    print("模型已保存")

writer.close()

7、模型的保存与读取

模型保存(method_2 保存的内存更小)

# save model of method 1 ==> structer + parameters
torch.save(vgg16_false, "save_model_method1.pth")

# save model of method 2 ==> parameters
torch.save(vgg16_false.state_dict(), "save_model_method2.pth")

模型加载

# load model of method 1
model1 = torch.load("save_model_method1.pth")

# load model of method 2
vgg16_false = torchvision.models.vgg16(pretrained=False)
vgg16_false.load_state_dict(torch.load("save_model_method2.pth"))

8、tensorboard使用

from torch.utils.tensorboard import SummaryWriter

writer = SummaryWriter("logs")

writer.add_image()
writer.add_images()
writer.add_scalar()

writer.close()

完整示例

import time

import torch.optim
import torchvision
from torch import nn
from torch.nn import CrossEntropyLoss
from torch.utils.tensorboard import SummaryWriter
from torch.nn import Linear, Flatten, MaxPool2d, Conv2d, Sequential
from torch.utils.data import DataLoader


# 定义训练设备
device = torch.device("cuda")

# 准备数据集
train_dataset = torchvision.datasets.CIFAR10(root="dataset_transforms", train=True, transform=torchvision.transforms.ToTensor(),
                                             download=True)
test_dataset = torchvision.datasets.CIFAR10(root="dataset_transforms", train=False, transform=torchvision.transforms.ToTensor(),
                                             download=True)

test_dataset_size = len(test_dataset)
# 加载数据集
train_loader = DataLoader(train_dataset, batch_size=64)
test_loader = DataLoader(test_dataset, batch_size=64)

# 搭建神经网络
class Module(nn.Module):
    def __init__(self):
        super(Module, self).__init__()
        self.model = Sequential(
            Conv2d(3, 32, 5, padding=2),
            MaxPool2d(2),
            Conv2d(32, 32, 5, padding=2),
            MaxPool2d(2),
            Conv2d(32, 64, 5, padding=2),
            MaxPool2d(2),
            Flatten(),
            Linear(1024, 64),
            Linear(64, 10)
        )

    def forward(self, input):
        output = self.model(input)
        return output
module = Module()
if torch.cuda.is_available():
    # 方法一
    # module = module.cuda()
    # 方法二
    module = module.to(device)

# 损失函数
loss_func = CrossEntropyLoss()
if torch.cuda.is_available():
    loss_func = loss_func.cuda()

# 优化器
learing_rate = 1e-2
optim = torch.optim.SGD(module.parameters(), lr=learing_rate)

# 记录训练次数
total_train_step = 0
# 记录测试次数
total_test_step = 0
# 训练轮数
epoch = 10


writer = SummaryWriter("logs_train_gpu")

for i in range(epoch):
    print("----------第{}轮训练----------".format(i + 1))
    start_time = time.time()
    # 模型训练步骤
    total_train_loss = 0
    for data in train_loader:
        imgs, targets = data
        if torch.cuda.is_available():
            imgs = imgs.cuda()
            targets = targets.cuda()
        outputs = module(imgs)
        loss = loss_func(outputs, targets)
        # 更新参数
        optim.zero_grad()
        loss.backward()
        optim.step()
        # 统计该轮训练的loss
        total_train_loss += loss
        # 训练次数
        total_train_step += 1

        if total_train_step % 100 == 0:
            print("第{}次的loss为 {}".format(total_train_step, loss.item()))
            end_time = time.time()
            print(end_time - start_time)
            writer.add_scalar("train_loss", loss.item(), total_train_step)

    # 模型测试步骤
    total_test_loss = 0
    total__accuracy = 0
    with torch.no_grad():
        for data in test_loader:
            imgs, targets = data
            if torch.cuda.is_available():
                imgs = imgs.cuda()
                targets = targets.cuda()
            outputs = module(imgs)
            loss = loss_func(outputs, targets)
            total_test_loss += loss
            accuracy = (outputs.argmax(1) == targets).sum()
            total__accuracy += accuracy

    print("测试集上的loss为 {}".format(total_test_loss))
    print("测试集上的准确率为 {}".format(total__accuracy / test_dataset_size))
    writer.add_scalar("test_loss", total_test_loss, total_test_step)
    writer.add_scalar("accuracy_rate", total__accuracy / test_dataset_size, total_test_step)
    total_test_step += 1

    torch.save(module, "module{}.pth".format(i))
    print("模型已保存")

writer.close()