pytorch训练模板

来源:http://worthpen.top/#/home/blog?blog=pot-blog36.md

引言

本项目实现了基于PyTorch Lightning的神经网络训练和测试管道。项目除了实现PyTorch Lightning的工作流外,还实现了通过任务池在训练过程中添加任务、k折交叉验证、将训练结果保存在.cvs中、接受随机种子进行恢复训练、将模型转换为.onnx和.tflite。

项目地址: https://github.com/shenyan233/machine_learning_template

使用方法

环境配置

python version:3.7-3.10

bash 复制代码
pip install -r requirements.txt

cuda and torch need to be installed by itself. Recommendation:

bash 复制代码
pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 torchaudio==0.12.1+cu113 --extra-index-url https://download.pytorch.org/whl/cu113

配置网络架构和数据集

整个项目目录结构如下:

bash 复制代码
.
├── dataset
│   └── {dataset_name}
│        ├── test
│        │   ├── image
│        │   │   └── *.png
│        │   └── label.txt
│        ├── train
│        │   ├── image
│        │   │   └── *.png
│        │   └── label.txt
│        └── ...
├── network
│   └── {network_name}
│         ├── network.py
│         └── ...
└── ...

上述表示的文件或文件夹需要预先配置。省略号所代表的文件或文件夹保持默认即可。

数据集需要调整为自己的数据集,图像(*.png)名称为对应label.txt中的行号。您可以在此处自由调整数据集的保存格式,但Dataloder和其他类需要在'/dataset/{dataset name}/init.py'中重写。

network.py包含要训练的网络架构, 可以改为其他名称, 但是需要同步修改__init__.py。

任务流的配置参数保存在tasks.json中。

设置训练参数

在./network/{network_name}/config.json中设置参数,参数包括model_name、dataset_path、stage、max_epoch、batch_size等。训练参数包括可选参数和必选参数, 具体内容可浏览main.py内的注释。参数stage为'fit'或'test',分别表示处于训练阶段或测试阶段。

开始训练或测试

在终端或cmd内执行:

python3 main.py

相关推荐
秃头佛爷7 分钟前
Python学习大纲总结及注意事项
开发语言·python·学习
昨日之日20061 小时前
Moonshine - 新型开源ASR(语音识别)模型,体积小,速度快,比OpenAI Whisper快五倍 本地一键整合包下载
人工智能·whisper·语音识别
浮生如梦_1 小时前
Halcon基于laws纹理特征的SVM分类
图像处理·人工智能·算法·支持向量机·计算机视觉·分类·视觉检测
深度学习lover1 小时前
<项目代码>YOLOv8 苹果腐烂识别<目标检测>
人工智能·python·yolo·目标检测·计算机视觉·苹果腐烂识别
热爱跑步的恒川2 小时前
【论文复现】基于图卷积网络的轻量化推荐模型
网络·人工智能·开源·aigc·ai编程
API快乐传递者2 小时前
淘宝反爬虫机制的主要手段有哪些?
爬虫·python
阡之尘埃4 小时前
Python数据分析案例61——信贷风控评分卡模型(A卡)(scorecardpy 全面解析)
人工智能·python·机器学习·数据分析·智能风控·信贷风控
孙同学要努力6 小时前
全连接神经网络案例——手写数字识别
人工智能·深度学习·神经网络
Eric.Lee20216 小时前
yolo v5 开源项目
人工智能·yolo·目标检测·计算机视觉
其实吧37 小时前
基于Matlab的图像融合研究设计
人工智能·计算机视觉·matlab