如何采用GPU训练_如何使用gpu训练

如何采用GPU训练

方法1：对网络模型，数据（数据、标注），损失函数调用.cuda()即可

import torch
import torchvision
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
import time
# 将model文件夹中有的东西都引入过来
# from model import *

# 准备数据集
train_data = torchvision.datasets.CIFAR10("data", train=True, transform=torchvision.transforms.ToTensor(),
                                          download=True)
test_data = torchvision.datasets.CIFAR10("data", train=False, transform=torchvision.transforms.ToTensor(),
                                         download=True)

# 看一下训练数据集和测试数据集有多少张 len-length 长度
train_data_size = len(train_data)
test_data_size = len(test_data)

# python中常用的写法：字符串格式化
print("训练数据集的长度为：{}".format(train_data_size))
print("测试数据集的长度为：{}".format(test_data_size))

# 用DataLoader加载数据集
train_dataloader = DataLoader(train_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)


# 搭建神经网络
class Peipei(nn.Module):
    def __init__(self):
        super(Peipei, self).__init__()
        self.model = nn.Sequential(
            nn.Conv2d(3, 32, 5, 1, 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


# 创建网络模型
peipei = Peipei()
if torch.cuda.is_available():
    peipei = peipei.cuda()

# 创建损失函数
loss_fn = nn.CrossEntropyLoss()
if torch.cuda.is_available():
    loss_fn = loss_fn.cuda()

# 定义优化器
# 1e-2 = 0.01
learning_rate = 1e-2
optimizer = torch.optim.SGD(peipei.parameters(), lr=learning_rate)

# 设置训练网络的一些参数
# 记录训练次数
total_train_step = 0
# 记录测试次数
total_test_step = 0
# 训练的轮数
epoch = 10

# 添加tensorboard
writer = SummaryWriter("logs_train")

start_time = time.time()

# i从0-9
for i in range(epoch):
    print("--------------------第{}轮训练开始--------------------".format(i + 1))

    # 训练步骤开始
    # 使模型进入训练状态，但只对特定层（Dropout,BatchNorm层）起作用
    peipei.train()
    for data in train_dataloader:
        imgs, targets = data
        if torch.cuda.is_available():
            imgs = imgs.cuda()
            targets = targets.cuda()
        outputs = peipei(imgs)
        # 计算损失函数
        loss = loss_fn(outputs, targets)

        # 优化器优化模型
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        total_train_step = total_train_step + 1
        if total_train_step % 100 == 0:
            end_time = time.time()
            print(end_time-start_time)
            print("训练次数：{},loss:{}".format(total_train_step, loss.item()))
            writer.add_scalar("train_loss", loss.item(), total_train_step)

    # 测试步骤开始
    # 使模型进入验证状态，但只对特定层（Dropout,BatchNorm层）起作用
    peipei.eval()
    total_test_loss = 0
    totel_accuracy = 0
    with torch.no_grad():
        for data in test_dataloader:
            imgs, targets = data
            if torch.cuda.is_available():
                imgs = imgs.cuda()
                targets = targets.cuda()
            outputs = peipei(imgs)
            loss = loss_fn(outputs, targets)
            # 计算整体测试集损失函数
            total_test_loss = total_test_loss + loss.item()
            # 计算整体正确率
            accuracy = (outputs.argmax(1) == targets).sum()
            totel_accuracy = totel_accuracy + accuracy

    print("整体测试集上的Loss：{}".format(total_test_loss))
    print("整体测试集上的正确率：{}".format(totel_accuracy / test_data_size))
    writer.add_scalar("test_loss", total_test_loss, total_test_step)
    writer.add_scalar("test_accuracy", totel_accuracy / test_data_size, total_test_step)
    total_test_step = total_test_step + 1

    # 对每轮训练完的模型保存
    torch.save(peipei, "peipei_{}.pth".format(i))
    torch.save(peipei.state_dict(), "peipei_{}.pth".format(i))

writer.close()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135

CPU

GPU

方法2

import torchvision
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
import time
# 将model文件夹中有的东西都引入过来
# from model import *

# 定义训练的设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 准备数据集
train_data = torchvision.datasets.CIFAR10("data", train=True, transform=torchvision.transforms.ToTensor(),
                                          download=True)
test_data = torchvision.datasets.CIFAR10("data", train=False, transform=torchvision.transforms.ToTensor(),
                                         download=True)

# 看一下训练数据集和测试数据集有多少张 len-length 长度
train_data_size = len(train_data)
test_data_size = len(test_data)

# python中常用的写法：字符串格式化
print("训练数据集的长度为：{}".format(train_data_size))
print("测试数据集的长度为：{}".format(test_data_size))

# 用DataLoader加载数据集
train_dataloader = DataLoader(train_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)


# 搭建神经网络
class Peipei(nn.Module):
    def __init__(self):
        super(Peipei, self).__init__()
        self.model = nn.Sequential(
            nn.Conv2d(3, 32, 5, 1, 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


# 创建网络模型
peipei = Peipei()
peipei = peipei.to(device)

# 创建损失函数
loss_fn = nn.CrossEntropyLoss()
loss_fn = loss_fn.to(device)

# 定义优化器
# 1e-2 = 0.01
learning_rate = 1e-2
optimizer = torch.optim.SGD(peipei.parameters(), lr=learning_rate)

# 设置训练网络的一些参数
# 记录训练次数
total_train_step = 0
# 记录测试次数
total_test_step = 0
# 训练的轮数
epoch = 10

# 添加tensorboard
writer = SummaryWriter("logs_train")

start_time = time.time()

# i从0-9
for i in range(epoch):
    print("--------------------第{}轮训练开始--------------------".format(i + 1))

    # 训练步骤开始
    # 使模型进入训练状态，但只对特定层（Dropout,BatchNorm层）起作用
    peipei.train()
    for data in train_dataloader:
        imgs, targets = data
        imgs = imgs.to(device)
        targets = targets.to(device)
        outputs = peipei(imgs)
        # 计算损失函数
        loss = loss_fn(outputs, targets)

        # 优化器优化模型
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        total_train_step = total_train_step + 1
        if total_train_step % 100 == 0:
            end_time = time.time()
            print(end_time - start_time)
            print("训练次数：{},loss:{}".format(total_train_step, loss.item()))
            writer.add_scalar("train_loss", loss.item(), total_train_step)

    # 测试步骤开始
    # 使模型进入验证状态，但只对特定层（Dropout,BatchNorm层）起作用
    peipei.eval()
    total_test_loss = 0
    totel_accuracy = 0
    with torch.no_grad():
        for data in test_dataloader:
            imgs, targets = data
            imgs = imgs.to(device)
            targets = targets.to(device)
            outputs = peipei(imgs)
            loss = loss_fn(outputs, targets)
            # 计算整体测试集损失函数
            total_test_loss = total_test_loss + loss.item()
            # 计算整体正确率
            accuracy = (outputs.argmax(1) == targets).sum()
            totel_accuracy = totel_accuracy + accuracy

    print("整体测试集上的Loss：{}".format(total_test_loss))
    print("整体测试集上的正确率：{}".format(totel_accuracy / test_data_size))
    writer.add_scalar("test_loss", total_test_loss, total_test_step)
    writer.add_scalar("test_accuracy", totel_accuracy / test_data_size, total_test_step)
    total_test_step = total_test_step + 1

    # 对每轮训练完的模型保存
    torch.save(peipei, "peipei_{}.pth".format(i))
    torch.save(peipei.state_dict(), "peipei_{}.pth".format(i))

writer.close()

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134

GPU

CPU

完整的模型验证（测试，demo）套路：利用已经训练好的模型，然后给它提供输入（对外应用）

GitHub常见代码