使用Pytorch-UNet训练自己的数据集

在开始之前，你需要安装Pytorch和Pytorch-UNet。你可以使用pip命令进行安装：

pip install torch torchvision
pip install git+https://github.com/milesial/Pytorch-UNet.git

第一步：制作数据集
将相同数量的图片和其掩码放置在data路径下的imgs和masks文件夹里，图片名需要和掩码标签相同。例如，如果你有一个名为’image1.jpg’的图片，你需要有一个名为’image1.png’的掩码。
第二步：修改utils/dataset.py（非必要，可跳过）
你需要将源代码下篮框标注的地方改成：

img_transform = transforms.Compose([
transforms.ToTensor(),
])
mask_transform = transforms.Compose([
transforms.ToTensor(),
])

第三步：修改train.py文件
train.py文件里有如图传参的地方，从上到下分别意思是：epoch，训练批次batch size，一次输入图片的数量learning rate，学习率load，加载预训练模型的路径scale，该值为0-1之间，设置越高，精度越高，但是资源占用越多validation，0-100之间，为验证集占数据中的百分比值。你可以根据实际情况进行修改。
第四步：开始训练模型
在终端中进入项目根目录，然后运行以下命令：

python train.py --epoch 10 --batch_size 4 --learning_rate 0.01 --load '' --scale 0.8 --validation 20

这个命令将会训练模型10个世代，每个世代使用4个样本进行训练，学习率为0.01，不加载预训练模型，scale设置为0.8，验证集占比为20%。你可以根据实际情况进行修改。
在训练过程中，你可以观察模型的训练曲线和损失值，以便进行调参和优化。如果你的模型出现了过拟合或者训练效果不好，你可以尝试使用其他的损失函数或者调整超参数。同时，你也可以尝试使用其他的优化器进行优化。在训练结束后，你可以使用测试脚本进行测试和评估模型的性能。测试脚本会自动保存预测结果和对应的真实标签。
总结：使用Pytorch-UNet训练自己的数据集需要准备数据集、修改代码和调整超参数等步骤。在训练过程中，你需要观察模型的训练曲线和损失值，以便进行调参和优化。如果你的模型出现了过拟合或者训练效果不好，你可以尝试使用其他的损失函数或者调整超参数。同时，你也可以尝试使用其他的优化器进行优化。在训练结束后，你可以使用测试脚本进行测试和评估模型的性能。