python version 3.8, torch 1.8.1, torchvision 0.9.1:
pip install torch==1.8.1 torchvision==0.9.1
数据文件夹结构如下:
datasets/
images/ # images
train/
img1.jpg
img2.jpg
.
.
.
val/
img1.jpg
img2.jpg
.
.
.
labels/ # masks
train/
img1.png
img2.png
.
.
.
val/
img1.png
img2.png
.
.
.
python train.py --input_size 224 224 --batch_size 32 --epochs 100 --nb_classes 2 --data_path ./datasets/ --output_dir ./output_dir
python eval.py --input_size 224 224 --batch_size 8 --weights ./output_dir/best.pth --data_path ./datasets/ --nb_classes 2
python predict.py --input_size 224 224 --weights ./output_dir/best.pth --image_path ./1.jpg --nb_classes 2
python export_onnx.py --input_size 224 224 --weights ./output_dir/best.pth --nb_classes 2
python -m onnxsim best.onnx best_sim.onnx