(args)
| 19 | |
| 20 | |
| 21 | def train(args): |
| 22 | # Read the config file # |
| 23 | with open(args.config_path, 'r') as file: |
| 24 | try: |
| 25 | config = yaml.safe_load(file) |
| 26 | except yaml.YAMLError as exc: |
| 27 | print(exc) |
| 28 | print(config) |
| 29 | ######################## |
| 30 | |
| 31 | dataset_config = config['dataset_params'] |
| 32 | train_config = config['train_params'] |
| 33 | |
| 34 | seed = train_config['seed'] |
| 35 | torch.manual_seed(seed) |
| 36 | np.random.seed(seed) |
| 37 | random.seed(seed) |
| 38 | if device == 'cuda': |
| 39 | torch.cuda.manual_seed_all(seed) |
| 40 | |
| 41 | voc = VOCDataset('train', |
| 42 | im_dir=dataset_config['im_train_path'], |
| 43 | ann_dir=dataset_config['ann_train_path']) |
| 44 | |
| 45 | train_dataset = DataLoader(voc, |
| 46 | batch_size=4, |
| 47 | shuffle=True, |
| 48 | num_workers=4, |
| 49 | collate_fn=collate_function) |
| 50 | |
| 51 | if args.use_resnet50_fpn: |
| 52 | faster_rcnn_model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True, |
| 53 | min_size=600, |
| 54 | max_size=1000, |
| 55 | ) |
| 56 | faster_rcnn_model.roi_heads.box_predictor = FastRCNNPredictor( |
| 57 | faster_rcnn_model.roi_heads.box_predictor.cls_score.in_features, |
| 58 | num_classes=21) |
| 59 | else: |
| 60 | backbone = torchvision.models.resnet34(pretrained=True, norm_layer=torchvision.ops.FrozenBatchNorm2d) |
| 61 | backbone = torch.nn.Sequential(*list(backbone.children())[:-3]) |
| 62 | backbone.out_channels = 256 |
| 63 | roi_align = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'], output_size=7, sampling_ratio=2) |
| 64 | rpn_anchor_generator = AnchorGenerator() |
| 65 | faster_rcnn_model = torchvision.models.detection.FasterRCNN(backbone, |
| 66 | num_classes=21, |
| 67 | min_size=600, |
| 68 | max_size=1000, |
| 69 | rpn_anchor_generator=rpn_anchor_generator, |
| 70 | rpn_pre_nms_top_n_train=12000, |
| 71 | rpn_pre_nms_top_n_test=6000, |
| 72 | box_batch_size_per_image=128, |
| 73 | rpn_post_nms_top_n_test=300 |
| 74 | ) |
| 75 | |
| 76 | faster_rcnn_model.train() |
| 77 | faster_rcnn_model.to(device) |
| 78 | if not os.path.exists(train_config['task_name']): |
no test coverage detected