pytorch模型搭建与训练

构建数据加载器

传入模型训练的数据常用DataLoader类封装(torch.utils.data.DataLoader),它能一次生成n条数据,即一个batch,n即设置的batch_size的大小,输出数据类型时Tensor,一般包含数据和对应的标签。同时可以设置是否乱序生成数据。如果数据集大小不能被batch_size整除,默认最后一个batch会取剩余的数据,如果想让剩余的省去来使每个batch数据量相同,可通过设置参数drop_last=True来实现。

1
data_loader = torch.utils.data.DataLoader(dataset=dataset,batch_size=batch_size, shuffle=True)

如果只需要训练数据,上面传入的dataset可以直接用list或者np.ndarray,如果要同时获得训练数据和标签,则需要对dataset用下面的方式构建一下:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
class GetLoader(torch.utils.data.Dataset):
# 初始化函数,得到数据
def __init__(self, data_root, data_label):
self.data = data_root
self.label = data_label

# index是根据batchsize划分数据后得到的索引,最后将data和对应的labels进行一起返回
def __getitem__(self, index):
data = self.data[index]
labels = self.label[index]
return data, labels

# 该函数返回数据大小长度,目的是DataLoader方便划分,如果不知道大小,DataLoader会一脸懵逼
def __len__(self):
return len(self.data)

参考链接:https://chenllliang.github.io/2020/02/04/dataloader/

搭建模型

  1. 构建模型类,继承nn.Module
  2. __init__ 方法里初始化父类init,然后定义模型需要的层,如线性、卷积
  3. 重写父类的forward方法,此方法传入模型输入,返回模型的输出。里面是利用init里定义的层来搭建网络。

模型类里还可以定义一些其他需要的运算函数,在forward里面可以调用。哪些运算定义在init里,哪些运算定义在forward里没有明确规定,通常的约定是网络的层写在init里,forward里用这些层来搭建网络。

调用很简单:y = model(x),因为是继承的类,forward函数会在父类里面自动被调用,我们无需关心细节。返回的y是Tensor 类型。

print(model) 可以来查看模型的结构。

一个CNN例子:

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
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
# 1 input image channel, 6 output channels, 5x5 square convolution
# kernel
self.conv1 = nn.Conv2d(1, 6, 5)
self.conv2 = nn.Conv2d(6, 16, 5)
# an affine operation: y = Wx + b
self.fc1 = nn.Linear(16 * 5 * 5, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10)
def forward(self, x):
# Max pooling over a (2, 2) window
x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
# If the size is a square you can only specify a single number
x = F.max_pool2d(F.relu(self.conv2(x)), 2)
x = x.view(-1, self.num_flat_features(x))
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x
def num_flat_features(self, x):
size = x.size()[1:] # all dimensions except the batch dimension
num_features = 1
for s in size:
num_features *= s
return num_features
1
2
3
4
5
6
7
8
9
net = Net()
print(net)
>>
Net(
(conv1): Conv2d(1, 6, kernel_size=(5, 5), stride=(1, 1))
(conv2): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))
(fc1): Linear(in_features=400, out_features=120, bias=True)
(fc2): Linear(in_features=120, out_features=84, bias=True) (fc3): Linear(in_features=84, out_features=10, bias=True)
)

训练模型

首先定义要训练的epoch数、batch_size大小、学习率learning_rate。

检查否有可用的GPU,有的话就用GPU训练:

1
2
3
4
5
6
if torch.cuda.is_available():
# model.cuda()
print('cuda is OK!')
model = model.to('cuda')
else:
print('cuda is NO!')

定义损失函数和优化器,这些在pytorch中都有实现,选择合适的来用就行,想要自己定义损失函数也可以。

训练的主要步骤如下:

1
2
3
4
5
6
7
8
9
10
11
for i in range(epochs):
model.train() # 开启训练模式
for x, label in dataloader: # 生成器,一个返回一个batch数据
optimizer.zero_grad() # 梯度清零
y = model(x) # 通过x求得y
loss = loss_function(y,label) # 求损失
loss.backward() # 梯度反向传播
optimizer.step() # 更新参数
##### 打印训练进度 ####
print(loss.item())
... # 打印其他信息

保存模型

1
torch.save(model.state_dict(), './model_.pth')

加载、使用模型

1
2
3
4
5
6
7
model = Net()
model.load_state_dict(torch.load('vae.pth'))
model.eval() # 开启训练模式,防止参数改变

data = np.loadtxt('data.csv', delimiter=',', dtype = np.float32) # 注意设定类型为np.float32
data = torch.Tensor(data) # 将np数组转换为Tensor
y_predict = model(data) # 预测

要注意的是从csv中读取时设定数据类型为np.float32,应为np读取默认为float64,而模型中参数为float32,不同类型无法运算。

同时需要将数据转换为torch.Tensor()类型才能运算。