(args)
| 14 | |
| 15 | |
| 16 | def train(args): |
| 17 | # Read the config file # |
| 18 | with open(args.config_path, 'r') as file: |
| 19 | try: |
| 20 | config = yaml.safe_load(file) |
| 21 | except yaml.YAMLError as exc: |
| 22 | print(exc) |
| 23 | print(config) |
| 24 | ######################## |
| 25 | |
| 26 | dataset_config = config['dataset_params'] |
| 27 | model_config = config['model_params'] |
| 28 | train_config = config['train_params'] |
| 29 | |
| 30 | seed = train_config['seed'] |
| 31 | torch.manual_seed(seed) |
| 32 | np.random.seed(seed) |
| 33 | random.seed(seed) |
| 34 | if device == 'cuda': |
| 35 | torch.cuda.manual_seed_all(seed) |
| 36 | |
| 37 | voc = VOCDataset('train', |
| 38 | im_dir=dataset_config['im_train_path'], |
| 39 | ann_dir=dataset_config['ann_train_path']) |
| 40 | train_dataset = DataLoader(voc, |
| 41 | batch_size=1, |
| 42 | shuffle=True, |
| 43 | num_workers=4) |
| 44 | |
| 45 | faster_rcnn_model = FasterRCNN(model_config, |
| 46 | num_classes=dataset_config['num_classes']) |
| 47 | faster_rcnn_model.train() |
| 48 | faster_rcnn_model.to(device) |
| 49 | |
| 50 | if not os.path.exists(train_config['task_name']): |
| 51 | os.mkdir(train_config['task_name']) |
| 52 | optimizer = torch.optim.SGD(lr=train_config['lr'], |
| 53 | params=filter(lambda p: p.requires_grad, |
| 54 | faster_rcnn_model.parameters()), |
| 55 | weight_decay=5E-4, |
| 56 | momentum=0.9) |
| 57 | scheduler = MultiStepLR(optimizer, milestones=train_config['lr_steps'], gamma=0.1) |
| 58 | |
| 59 | acc_steps = train_config['acc_steps'] |
| 60 | num_epochs = train_config['num_epochs'] |
| 61 | step_count = 1 |
| 62 | |
| 63 | for i in range(num_epochs): |
| 64 | rpn_classification_losses = [] |
| 65 | rpn_localization_losses = [] |
| 66 | frcnn_classification_losses = [] |
| 67 | frcnn_localization_losses = [] |
| 68 | optimizer.zero_grad() |
| 69 | |
| 70 | for im, target, fname in tqdm(train_dataset): |
| 71 | im = im.float().to(device) |
| 72 | target['bboxes'] = target['bboxes'].float().to(device) |
| 73 | target['labels'] = target['labels'].long().to(device) |
no test coverage detected