使用torch定义自己的数据集与网络结构
在学习完如何配置torch开发环境,如何使用框架提供的数据集,以及如何简单搭建训练流程之后,我们可以开始使用torch定义自己的数据集与网络结构。
一、使用torch定义数据集
torch提供了torch.utils.data.Dataset类,该类是所有数据集的基类,它实现了序列协议,即__len__和__getitem__方法,前者返回数据集的大小,后者返回第i个样本。
from torch.utils.data import Dataset
class MyDataset(Dataset):
def __init__(self, data_path):
...
def __len__(self):
...
def __getitem__(self, index):
...
- 在__init__方法中,我们需要从文件中加载数据,但是一般的做法是保存每条数据的的路径,因为对于大型数据集,将其直接加载到内存可能会造成大量的内存浪费
- 在__len__方法中,我们只需要返回数据集的大小
- 在__getitem__方法中,我们需要返回第index个样本和其对应的标签,同时,可能还需要进行一些数据增强的操作
这样就实现了自定义数据集。使用torch.utils.data.DataLoader就可以从我们的数据集中加载数据,并自动进行一些数据增强,采样等操作。
from torch.utils.data import DataLoader
from torch.utils.data import Dataset
class MyDataset(Dataset):
...
data_path='data_path'
mydataset=MyDataset(data_path)
dataloader=DataLoader(mydataset,batch_size=32,shuffle=True)
for index,batch in enumerate(dataloader):
...
dataloader的参数包括batch_size,shuffle,num_workers等等,其中batch_size表示每次训练的样本数,shuffle表示是否对数据进行打乱,num_workers表示使用多少个进程来读取数据。
也就是说,我们不需要关注批次数据的采样和加载,dataloader会自动帮我们完成。我们只需要实现数据集的__len__和__getitem__方法,然后使用dataloader即可。
二、使用torch定义网络结构
torch提供了torch.nn.Module类,该类是所有torch中神经网络的基类,通过继承该类,并实现forward方法,就可以定义自己的网络结构。
import torch
from torch import nn
class MyNet(nn.Module):
def __init__(self,input_dim,output_dim):
# input_dim表示输入的维度,output_dim表示输出的维度,可以根据需要自行定义其他参数
# 也可以直接 硬编码 ,将定义网络的相关参数写死
# 在初始化方法中,我们需要定义网络结构,包括网络中的层,以及层的参数
super(MyNet,self).__init__()
# 全连接层
self.fc1=nn.Linear(input_dim,64)
self.fc2=nn.Linear(64,32)
self.fc3=nn.Linear(32,output_dim)
# 激活函数
self.relu=nn.ReLU()
def forward(self,x):
x=self.fc1(x)
x=self.relu(x)
x=self.fc2(x)
x=self.relu(x)
x=self.fc3(x)
return x
这样我们就实现了一个简单的全连接神经网络。
torch提供了相当多的算子,我们可以根据需要使用它们来定义自己的网络结构。这些算子本质上也是一个继承了torch.nn.Module类的类,因此我们可以直接使用它们来定义自己的网络结构。
例如
- torch.nn.Linear 全连接层
- torch.nn.ReLU 激活函数
- torch.nn.Conv2d 卷积层
- torch.nn.MaxPool2d 最大池化层
- torch.nn.BatchNorm2d 批归一化层
- torch.nn.Dropout 随机失活层
- torch.nn.Sequential 顺序容器,顺序执行其中的算子,简化网络结构
各种算子可以参考官方文档
torch.nn
在forward方法中,我们需要组织网络结构,将在初始化方法中定义好的网络结构按照我们预想的顺序组织起来,并返回结果。
三、使用torch定义损失函数
损失函数实际上也是一个算子,因此我们可以通过继承torch.nn.Moudle类,用定义网络结构的方法来定义损失函数,不过这时forward方法需要两个参数,即预测结果和真实标签。然后我们在forward方法中按照某种方法计算两者之间的差异,这个差异就是我们需要的损失。
torch已经提供了丰富的损失函数,我们可以直接使用它们
- torch.nn.CrossEntropyLoss 交叉熵损失,用于多分类问题
- torch.nn.MSELoss 均方差损失,用于回归问题
- torch.nn.BCELoss 二元交叉熵损失,用于二分类问题
各种损失函数可以参考官方文档
torch.nn
也可以根据需要自定义损失函数。
例如 目标检测 yolo 的损失函数,是多种差异计算结果的组合,torch没有对应的封装,需要自己实现。
四、基于自己的数据集和网络结构训练模型
定义好数据集和网络结构之后,就可以使用torch.optim.Optimizer类来定义优化器,并使用torch.nn.CrossEntropyLoss类来定义损失函数,然后搭建训练流程。
import torch
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.data import Dataset
class MyDataset(Dataset):
...
class MyNet(nn.Module):
...
place='cuda' if torch.cuda.is_available() else 'cpu' # 选择使用cpu还是gpu
data_path='data_path'
mydataset=MyDataset(data_path)
dataloader=DataLoader(mydataset,batch_size=32,shuffle=True)
model=MyNet(input_dim,output_dim)
model=model.to(place)
criterion=nn.CrossEntropyLoss()
optimizer=torch.optim.SGD(model.parameters(),lr=0.01)
for epoch in range(100):
for index,batch in enumerate(dataloader):
input,label=batch
input=input.to(place)
label=label.to(place)
# 前向传播(forward)
output=model(input)
# 计算损失
loss=criterion(output,label)
# 梯度清零
optimizer.zero_grad()
# 反向传播(框架会自动计算梯度,无需手动计算)
loss.backward()
# 更新参数
optimizer.step()
print('epoch:{},index:{},loss:{}'.format(epoch,index,loss.item()))
torch.save(model.state_dict(),'model.pth')
上述代码实现了一个简单的训练流程,其中包含数据集的读取、网络结构的定义、损失函数的定义、优化器的定义和训练流程,可以根据需要,添加训练过程中需要执行的操作。
例如:
- 一般在每个epoch之后,可以保存模型的参数,以便在训练中断后继续训练。
- 可以在训练过程中,使用验证集来评估模型的性能,以确定是否需要停止训练。
- 可以在训练过程中,使用测试集来评估模型的性能,以确定模型的泛化能力。
- 可以在训练过程中,使用tensorboard来可视化训练过程,以帮助调试和优化模型。
接下来我会用一个完整的项目作为示范
torch使用自己的数据集与网络模型
DataSet.py
我从aistudio上下载了 马或人类数据集 用于演示
结构如下

因此我们需要根据数据集的结构编写合适的数据集类:
- 使用 os.listdir() 方法读取文件夹中的文件名,并使用 os.path.join() 方法拼接文件路径。
- 加载图片,并使用torchvision.transforms进行图像增强,并缩放到网络需要的大小
- 返回图片和标签
Net.py
然后定义网络模型,这里为了获得更好的训练效果,我使用了迁移学习的方法,使用预训练的resnet50作为骨干网络,添加了并添加全连接层用于分类,正如之前说的,torch中的算子都是继承了torch.nn.Module类的,因此我们可以直接使用它们来定义自己的网络结构,即使是完整的网络模型,也同样可以当作算子来使用。
train.py
在这个文件中定义训练流程,包括训练,测试,推理
我已经简单的进行了训练,可以在这个文件的入口处注释掉train()从而跳过训练,直接进行推理
因为案例比较简单,又使用了预训练的resnet50作为骨干网络,因此在训练集和测试集上的准确率几乎都能达到100%
运行结果
为了检测模型的实际性能,我从从百度图库中随机下载了人类和马的图片用于推理,最终都得到了正确的分类结果。可以尝试输入不同图片,查看模型的分类效果。
总结
如果我们想要在torch中使用自定义的算子和数据集,我们只需要继承对应的基类,并重写对应的方法,那么就能像使用torch自带的模块一样使用我们自定义的模块,保证接口的一致性,让开发者将更多的精力用于算法的实现
关于迁移学习
简单来说,迁移学习就是使用别人训练好的模型在我们自己的数据集上进一步微调。
在训练一个网络模型时,通常需要大量的数据,数据集的规模越大,训练所需的时间越多,精度越高。但是,我们自己的数据集往往没有那么大,因此,我们需要在预训练的模型基础上进行微调。预训练模型已经具有了优秀的特征提取能力,所以在迁移学习中主要训练的就是输出层,这样就极大的缩短了训练所需的时间,提高了模型的精度。
2024/7/23

被折叠的 条评论
为什么被折叠?



