train_demo.py 1.3 KB

123456789101112131415161718192021222324252627282930
  1. import torch
  2. from models.line_detect.line_detect import linedetect_newresnet18fpn, linedetect_resnet50_fpn, linedetect_resnet18_fpn, \
  3. linedetect_newresnet50fpn, linedetect_maxvitfpn, linedetect_high_maxvitfpn, linedetect_swin_transformer_fpn, \
  4. linedetect_newresnet101fpn, linedetect_newresnet152fpn, linedetect_efficientnet
  5. from models.line_net.trainer import Trainer
  6. device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
  7. if __name__ == '__main__':
  8. # model = LineNet('line_net.yaml')
  9. # model=linedetect_resnet50_fpn()
  10. # model = linedetect_resnet50_fpn()
  11. # model=get_line_net_convnext_fpn(num_classes=2).to(device)
  12. # model=linenet_newresnet50fpn()
  13. # model = lineDetect_resnet18_fpn()
  14. # model=linedetect_resnet18_fpn()
  15. # model=linedetect_newresnet18fpn(num_points=4)
  16. # model=linedetect_newresnet50fpn(num_points=4)
  17. # model = linedetect_newresnet101fpn(num_points=4)
  18. # model = linedetect_newresnet152fpn(num_points=4)
  19. # model.load_weights(save_path=r'/home/admin/projects/MultiVisionModels/models/line_detect/train_results/20250711_114046/weights/best_val.pth')
  20. model=linedetect_maxvitfpn()
  21. # model=linedetect_efficientnet(name='efficientnet_v2_l')
  22. # model=linedetect_high_maxvitfpn()
  23. # model=linedetect_swin_transformer_fpn(type='t')
  24. model.start_train(cfg='train.yaml')