AllentDan/LibtorchSegmentation

cuda

yuanhs1996 opened this issue · 2 comments

yuanhs1996 commented

Can training be accelerated using CUDA? If so, how can it be done?

AllentDan commented

Yes, examples are in the readme.

Segmentor<FPN> segmentor;
segmentor.Initialize(0/*gpu id, -1 for cpu*/,
                    512/*resize width*/,
                    512/*resize height*/,
                    {"background","person"}/*class name dict, background included*/,
                    "resnet34"/*backbone name*/,
                    "your path to resnet34.pt");
segmentor.Train(0.0003/*initial leaning rate*/,
                300/*training epochs*/,
                4/*batch size*/,
                "your path to voc_person_seg",
                ".jpg"/*image type*/,
                "your path to save segmentor.pt");
yuanhs1996 commented