webstudio/web_p/train_Dope.py
2024-05-02 17:36:44 +03:00

30 lines
1.5 KiB
Python

"""
train_Dope
Общая задача: оценка позиции объекта (Pose estimation)
Реализуемая функция: обучение нейросетевой модели DOPE по заданному BOP-датасету
python3 $PYTHON_EDUCATION --path /Users/idontsudo/webservice/server/build/public/7065d6b6-c8a3-48c5-9679-bb8f3a690296 \
--name test1234 --datasetName 32123213
25.04.2024 @shalenikol release 0.1
"""
import os
def train_Dope_i(path:str, wname:str, dname:str, outpath:str, epochs:int):
results = f"torchrun --nproc_per_node=1 train.py --local_rank 0 --data {os.path.join(path,dname)} --object fork" \
+ f" -e {epochs} --batchsize 16 --exts jpg --imagesize 640 --pretrained" \
+ " --net_path /home/shalenikol/fork_work/dope_training/output/weights_2996/net_epoch_47.pth"
print(results)
import argparse
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--path", required=True, help="Path for dataset")
parser.add_argument("--name", required=True, help="String with result weights name")
parser.add_argument("--datasetName", required=True, help="String with dataset name")
parser.add_argument("--outpath", default="weights", help="Output path for weights")
parser.add_argument("--epoch", default=3, help="How many training epochs")
args = parser.parse_args()
train_Dope_i(args.path, args.name, args.datasetName, args.outpath, args.epoch)