(config, logger, neptune)
| 35 | |
| 36 | |
| 37 | def train(config, logger, neptune): |
| 38 | |
| 39 | #! Define Model |
| 40 | lcd, gpu_conf = set_lcd_model(config, logger, neptune) |
| 41 | [_, device, gpu_ids] = gpu_conf |
| 42 | |
| 43 | #! Define Dataloader |
| 44 | train_loader = make_data_loader(config, gpu_ids, is_train=True) |
| 45 | if config.DATA.DATASET_NAME in ["PITT", "Campus"]: |
| 46 | eval_loader = make_data_loader(config, gpu_ids, is_train=False) |
| 47 | else: |
| 48 | eval_loader = [] |
| 49 | log_print("train batch with {}, eval batch with {}".format( |
| 50 | len(train_loader), len(eval_loader)), 'g') |
| 51 | |
| 52 | #! Define Tester |
| 53 | if config.DATA.DATASET_NAME == "PITT": |
| 54 | tester = EvaluationPitts(config, lcd, device) |
| 55 | |
| 56 | #! Main loop |
| 57 | prev_time = time.time() |
| 58 | best_recall = 0 |
| 59 | for epoch in range(lcd.epoch, config.TRAINING.EPOCH+1): |
| 60 | |
| 61 | log_print("Train epoch {}".format(epoch), "g") |
| 62 | |
| 63 | #! Do Training |
| 64 | for i, batch in enumerate(train_loader): |
| 65 | |
| 66 | # * Determine approximate time left |
| 67 | batches_done = epoch * len(train_loader) + i |
| 68 | batches_left = config.TRAINING.EPOCH * \ |
| 69 | len(train_loader) - batches_done |
| 70 | time_left = datetime.timedelta( |
| 71 | seconds = batches_left * (time.time() - prev_time)) |
| 72 | prev_time = time.time() |
| 73 | data = [x.to(device, dtype=torch.float) for x in batch] |
| 74 | loss_lidar = lcd.train_lcd(data) |
| 75 | sys.stdout.write( |
| 76 | "\r[Epoch %d/%d] [Batch %d/%d] [LiDAR loss: %f], ETA: %s" |
| 77 | % (epoch, config.TRAINING.EPOCH, i, len(train_loader), loss_lidar, time_left) |
| 78 | ) |
| 79 | |
| 80 | #! Do Evaluation |
| 81 | if config.DATA.DATASET_NAME in ["PITT", "Campus"]: |
| 82 | for i, batch in tqdm(enumerate(eval_loader), total=len(eval_loader)): |
| 83 | data = [x.to(device, dtype=torch.float) for x in batch] |
| 84 | lcd.eval_lcd(data) |
| 85 | |
| 86 | #! Do Test |
| 87 | if config.DATA.DATASET_NAME in ["PITT", "Campus"]: |
| 88 | test_stats = [] |
| 89 | for traj_num in config.DATA.TEST_LIST: |
| 90 | if config.DATA.DATASET_NAME == "PITT": |
| 91 | recall, _, _ = tester.get_features_recall(traj_num, 2, 0, 0) |
| 92 | else: |
| 93 | recall, _, _ = tester.infer(traj_num) |
| 94 | test_stats.append(recall) |
no test coverage detected