UNet跑自己的VOC格式数据集——超好用!Pytorch版本(自用笔记)

源码:https://github.com/bubbliiiing/unet-pytorch
参考的blog:https://blog.csdn.net/weixin_44791964/article/details/108866828
这个源码里的注释超级超级详细,我哭死。。。。

默认已经准备好了自己的VOC格式的数据集。

用train.py训练,get_miou.py预测就行了。
注意,如果要指定GPU的id的话,可以在import torch前面加上这句,数字就是id

os.environ["CUDA_VISIBLE_DEVICES"] = "1"
import torch

训练:
修改train.py文件:
是否使用Cuda,num_classes类别数,主干网络,model_path改为自己在readme里面下载好的权值文件路径,注意修改save_dir路径自己要在logs下建一个当前项目的子文件夹(要不然训练多了会混起来),数据集路径,dice_loss和focal_loss按照注释修改,还有一些训练参数自己按需修改。
预测mask:
因为我要预测多张图,predict.py里面还得自己改代码写遍历,后面发现get_miou.py里有现成的。设置miou_mode为1代表仅仅获得预测结果。
预测一般要改的:
get_miou.py里:miou_mode、分类数、种类、VOC路径、txt文件路径、输出mask的路径。
unet.py里:model_path指向./logs文件夹下的权值文件,num_classes,具体要看自己要更改哪些参数

  • 1
    点赞
  • 6
    收藏
    觉得还不错? 一键收藏
  • 4
    评论
对于使用 PyTorch 训练自己的数据集,你可以按照以下步骤进行: 1. 准备数据集:将你的数据集划分为训练集和验证集,并组织成 PyTorch 的 Dataset 类的形式。Dataset 类需要实现 `__len__()` 和 `__getitem__()` 方法,用于返回数据集大小和获取样本。 2. 数据预处理:根据你的任务需求,对图像进行必要的预处理操作,例如缩放、裁剪、归一化等。你可以使用 PyTorch 提供的图像处理工具包 torchvision 来方便地完成这些操作。 3. 定义网络模型:使用 PyTorch 构建 UNet 模型。你可以自己实现模型结构,也可以使用现有的开源实现。 4. 定义损失函数:根据你的任务类型,选择适当的损失函数。例如,对于图像分割任务,你可以使用交叉熵损失函数或 Dice Loss。 5. 定义优化器:选择合适的优化器来更新模型的参数。常用的优化器包括 Adam、SGD 等,你可以根据自己的需求进行选择。 6. 训练模型:使用 DataLoader 来加载数据,将数据输入到网络中进行训练。在每个 epoch 结束后,计算损失函数并进行反向传播更新模型参数。 7. 评估模型:使用验证集对训练的模型进行评估,计算预测结果的准确率、召回率、F1 值等指标。 8. 预测新数据:使用训练好的模型对新数据进行预测。将新数据输入到模型中,得到预测结果。 这些是基本的步骤,你可以根据自己的具体情况进行调整和扩展。希望这些对你有所帮助!

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 4
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值