def main(): model = Yolov1(split_size=7, num_boxes=2, num_classes=20).to(DEVICE) optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY) loss_fn = YoloLoss() train_dataset = VOCDataset("data/train.csv", transform=transform, img_dir=IMG_DIR, label_dir=LABEL_DIR) test_dataset = VOCDataset("data/test.csv", transform=transform, img_dir=IMG_DIR, label_dir=LABEL_DIR) train_loader=DataLoader(dataset=train_dataset, batch_size=BATCH_SIZE, num_workers=1, pin_memory=PIN_MEMORY, shuffle=True,drop_last=True) test_loader=DataLoader(dataset=test_dataset, batch_size=BATCH_SIZE, num_workers=1, pin_memory=PIN_MEMORY, shuffle=True,drop_last=True) for epoch in range(EPOCHS): pred_boxes, target_boxes = get_bboxes(train_loader, model, iou_threshold=0.5, threshold=0.4) mAP = mean_average_precision(pred_boxes, target_boxes, iou_threshold=0.5) print(f"Train mAP:{mAP}") train_fn(train_loader, model, optimizer, loss_fn) if epoch > 99: for x, y in test_loader: x = x.to(DEVICE) for idx in range(16): bboxes = cellboxes_to_boxes(model(x)) bboxes = non_max_suppression(bboxes[idx], iou_threshold=0.5, threshold=0.4) plot_image(x[idx].permute(1,2,0).to("cpu"), bboxes) if __name__ == "__main__": main()
def main(): yolo = DarkNet19() batch_data = pascal.generate_batch_data() yololoss = YoloLoss() optimizer = tf.optimizers.Adam(learning_rate=0.01) for epoch in range(EPOCHS): i = 0 for batch in batch_data.as_numpy_iterator(): batch_names, batch_boxs = pascal.parse_batch_set(batch) train_image = pascal.get_train_images(batch_names) train_labels = pascal.get_train_labels(batch_boxs) with tf.GradientTape() as tape: y_pred = yolo(train_image) loss = yololoss(y_true=train_labels, y_pred=y_pred) grads = tape.gradient(loss, yolo.trainable_variables) optimizer.apply_gradients(zip(grads, yolo.trainable_weights)) i = i + 1 print("epoch {} batch_index={} loss={}".format( epoch, i, loss.numpy())) if epoch % 1 == 0: print("epoch is {} loss is {}".format(epoch, loss))
def main(): model = YOLOv3(num_classes=config.NUM_CLASSES).to(config.DEVICE) optimizer = optim.Adam(model.parameters(), lr=config.LEARNING_RATE, weight_decay=config.WEIGHT_DECAY) loss_fn = YoloLoss() scaler = torch.cuda.amp.GradScaler() train_loader, test_loader, train_eval_loader = get_loaders( train_csv_path=config.DATASET + "/train.csv", test_csv_path=config.DATASET + "/test.csv") if config.LOAD_MODEL: load_checkpoint(config.CHECKPOINT_FILE, model, optimizer, config.LEARNING_RATE) print("loaded model") scaled_anchors = (torch.tensor(config.ANCHORS) * torch.tensor( config.S).unsqueeze(1).unsqueeze(1).repeat(1, 3, 2)).to(config.DEVICE) for epoch in range(config.NUM_EPOCHS): print(f"\nEpoch [{epoch}]") #plot_couple_examples(model, test_loader, 0.6, 0.5, scaled_anchors) train_fn(train_loader, model, optimizer, loss_fn, scaler, scaled_anchors) #if config.SAVE_MODEL: # save_checkpoint(model, optimizer, filename=f"checkpoint.pth.tar") #print("On Train Eval loader:") #print("On Train loader:") #check_class_accuracy(model, train_loader, threshold=config.CONF_THRESHOLD) if epoch > 0 and epoch % 4 == 0: valid_fn(train_eval_loader, model, loss_fn, scaled_anchors) check_class_accuracy(model, test_loader, threshold=config.CONF_THRESHOLD) pred_boxes, true_boxes = get_evaluation_bboxes( test_loader, model, iou_threshold=config.NMS_IOU_THRESH, anchors=config.ANCHORS, threshold=config.CONF_THRESHOLD, ) mapval = mean_average_precision( pred_boxes, true_boxes, iou_threshold=config.MAP_IOU_THRESH, box_format="midpoint", num_classes=config.NUM_CLASSES, ) print(f"MAP: {mapval.item()}") print("\nnResuming Training\n") model.train()
def main(): model = Yolov1(split_size=7, num_boxes=2, num_classes=20).to(DEVICE) optimizer = optim.Adam( model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY ) loss_fn = YoloLoss() if LOAD_MODEL: load_checkpoint(torch.load(LOAD_MODEL_FILE), model, optimizer) train_dataset = VOCDataset( "data/100examples.csv", transform=transform, img_dir=IMG_DIR, label_dir=LABEL_DIR ) test_dataset = VOCDataset( "data/test.csv", transform=transform, img_dir=IMG_DIR, label_dir=LABEL_DIR ) train_loader = DataLoader( dataset=train_dataset, batch_size=BATCH_SIZE, num_workers=NUM_WORKERS, pin_memory=PIN_MEMORY, shuffle=True, drop_last=True ) test_loader = DataLoader( dataset=test_dataset, batch_size=BATCH_SIZE, num_workers=NUM_WORKERS, pin_memory=PIN_MEMORY, shuffle=True, drop_last=True ) for epoch in range(EPOCHS): pred_boxes, target_boxes = get_bboxes( train_loader, model, iou_threshold=0.5, threshold=0.4 ) mean_avg_prec = mean_average_precision( pred_boxes, target_boxes, iou_threshold=0.5, box_format="midpoint" ) print(f"Train mAP in {epoch}: {mean_avg_prec}") train_fn(train_loader, model, optimizer, loss_fn)
def main(): model = YOLOv3(num_classes=config.NUM_CLASSES).to(config.DEVICE) optimizer = optim.Adam( model.parameters(), lr=config.LEARNING_RATE, weight_decay=config.WEIGHT_DECAY ) loss_fn = YoloLoss() scaler = torch.cuda.amp.GradScaler() train_loader, test_loader, train_eval_loader = get_loaders( train_csv_path=config.DATASET + "/train.csv", test_csv_path=config.DATASET + "/test.csv" ) if config.LOAD_MODEL: load_checkpoint( config.CHECKPOINT_FILE, model, optimizer, config.LEARNING_RATE ) scaled_anchors = ( torch.tensor(config.ANCHORS) * torch.tensor(config.S).unsqueeze(1).unsqueeze(1).repeat(1, 3, 2) ).to(config.DEVICE) for epoch in range(config.NUM_EPOCHS): train_fn(train_loader, model, optimizer, loss_fn, scaler, scaled_anchors) if epoch > 0 and epoch % 3 == 0: check_class_accuracy(model, test_loader, threshold=config.CONF_THRESHOLD) pred_boxes, true_boxes = get_evaluation_bboxes( test_loader, model, iou_threshold=config.NMS_IOU_THRESH, anchors=config.ANCHORS, threshold=config.CONF_THRESHOLD, ) mapval = mean_average_precision( pred_boxes, true_boxes, iou_threshold=config.MAP_IOU_THRESH, box_format="midpoint", num_classes=config.NUM_CLASSES, ) print(f"MAP: {mapval.item()}") model.train() if epoch > 99: for x, y in train_loader: x = x.to(DEVICE) for idx in range(8): bboxes = cellboxes_to_boxes(model(x)) bboxes = non_max_suppression(bboxes[idx], iou_threshold=0.5, threshold=0.4) plot_image(x[idx].permute(1,2,0).to("cpu"), bboxes, idx) if __name__ == "__main__": main()
def main(): # Data loading train_dataset = CustomDataset(root=config.root_train, annFile=config.annFile_train, transforms=config.train_transforms, catagory=config.CATEGORY_FILTER) val_dataset = CustomDataset(root=config.root_train, annFile=config.annFile_train, transforms=config.val_transforms, catagory=config.CATEGORY_FILTER) train_loader = DataLoader(dataset=train_dataset, batch_size=16, num_workers=4, pin_memory=True, shuffle=False, drop_last=True) val_loader = DataLoader(dataset=val_dataset, batch_size=16, num_workers=4, pin_memory=True, shuffle=False, drop_last=True) # Model model = YoloV3(num_classes=config.C).to(device=config.DEVICE) optimizer = optim.Adam(model.parameters(), lr=config.LEARNING_RATE, weight_decay=config.WEIGHT_DECAY) loss_function = YoloLoss() scalar = torch.cuda.amp.GradScaler() # Miscellaneous scaled_anchors = (torch.tensor(config.anchors) * torch.tensor( config.Scale).unsqueeze(1).unsqueeze(1).repeat(1, 3, 2)).to( config.DEVICE) #writer = SummaryWriter() current_time = time.time() print("Train loader length:", len(train_loader)) # Training loop model.train() for cycle, (x, y) in enumerate(train_loader): print("Current cycle:", cycle) delta_time, current_time = time_function(current_time)
def main(): model = YOLOv3(num_classes=config.NUM_CLASSES).to(config.DEVICE, dtype=torch.float) optimizer = optim.Adam(model.parameters(), lr=config.LEARNING_RATE, weight_decay=config.WEIGHT_DECAY) loss_fn = YoloLoss().to(config.DEVICE) scaler = torch.cuda.amp.GradScaler() dataset = YoloDataset(table=table, anchors=anchors, transform=config.transform) train_loader = DataLoader(dataset, batch_size=2, shuffle=True) for epoch in (range(20)): train_fn(train_loader, model, optimizer, loss_fn, scaler, anchors)
def main(): model = Yolov1(split_size=7, num_boxes=2, num_classes=20).to(DEVICE) optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY) loss_fn = YoloLoss() if LOAD_MODEL: pass train_transform = transforms.Compose( [transforms.Resize(size=(448, 448)), transforms.ToTensor()]) train_dataset = VOCDataset(csv_file='', img_root=IMG_DIR, S=7, B=2, C=20, transform=train_transform) test_transform = transforms.Compose( [transforms.Resize(size=(448, 448)), transforms.ToTensor()]) test_dataset = VOCDataset(csv_file='', img_root=IMG_DIR, transform=test_transform) train_loader = DataLoader( dataset=train_dataset, batch_size=BATCH_SIZE, num_workers=NUM_WORKERS, pin_memory=PIN_MEMORY, shuffle=True, drop_last=True, ) test_loader = DataLoader(dataset=test_dataset, batch_size=BATCH_SIZE, num_workers=NUM_WORKERS, pin_memory=PIN_MEMORY, shuffle=True, drop_last=True) for epoch in range(EPOCHS): train_fn(train_loader, model, optimizer, loss_fn)
def __init__(self, config, dev='0'): super(Trainer, self).__init__() with open(config, 'r') as f: config = json.load(f) self.iterations = config['train']['iterations'] # self.model = PneuNet(config['model']['img_shape'], # config['model']['classes']) # self.objective = nn.CrossEntropyLoss() self.model = PneuYoloNet(config['model']['img_shape']) self.objective = YoloLoss(config['model']['img_shape'][1]) lr = config['train']['learning_rate'] self.optimizer = optim.SGD(self.model.parameters(), lr=lr) self.losses = [] self.vallosses = [] if torch.cuda.is_available(): torch.cuda.set_device(int(dev)) self.model.cuda() self.dtype = torch.cuda.FloatTensor print("Using GPU") else: self.dtype = torch.FloatTensor print("No GPU detected") batch_size = config['train']['batch_size'] trainset = PneuDataset(config['data']['train'], config['data']['image_path'], config['model']['img_shape'][1]) self.trainloader = DataLoader(trainset, batch_size=batch_size, pin_memory=True) self.valset = PneuDataset(config['data']['valid'], config['data']['image_path'], config['model']['img_shape'][1]) self.valloader = DataLoader(self.valset, batch_size=batch_size, pin_memory=True) self.write_interval = config['model']['write_interval'] self.train_info_path = config['model']['train_info_path'] self.model_path = config['model']['model_save_path'].split('.pt')[0] self.graph_path = config['model']['graph_save_path'].split('.png')[0]
def main(): model = YOLOv1(split_size=7, num_boxes=2, num_classes=20).to(device) optimizer = optim.Adam(model.parameters(), lr=learning_rate, weight_decay=wd) loss_fn = YoloLoss() if load_model: load_checkpoint(torch.load(load_model_file), model, optimizer) train_dataset = VOCDataset("data/8examples.csv", transform=transform, img_dir=img_dir, label_dir=label_dir) test_dataset = VOCDataset("data/test.csv", transform=transform, img_dir=img_dir, label_dir=label_dir) train_loader = DataLoader(dataset=train_dataset, batch_size=bs, num_workers=num_workers, pin_memory=pin_mem, shuffle=True, drop_last=False) test_loader = DataLoader(dataset=test_dataset, batch_size=bs, num_workers=num_workers, pin_memory=pin_mem, shuffle=True, drop_last=True) for epoch in range(epochs): pred_boxes, target_boxes = get_bboxes(train_loader, model, iou_threshold=0.5, threshold=0.4) mean_avg_prec = mean_average_precision(pred_boxes, target_boxes, iou_threshold=0.5, box_format="midpoint") print(f"Train mAP: {mean_avg_prec}") train_fn(train_loader, model, optimizer, loss_fn)
def test_yolo_loss_0(self): # IF predictions = torch.zeros((16, 3, 13, 13, 25)) predictions[0, 0, 0, 0, 0:5] = torch.tensor([1.0, 0.5, 0.5, 0.5, 0.5]) predictions[0, 0, 0, 0, 6] = torch.tensor([1.0]) target = torch.zeros((16, 3, 13, 13, 6)) target[0, 0, 0, 0, ...] = torch.tensor([1.0, 0.5, 0.5, 0.5, 0.5, 1.0]) anchors = [[(0.28, 0.22), (0.38, 0.48), (0.9, 0.78)], [(0.07, 0.15), (0.15, 0.11), (0.14, 0.29)], [(0.02, 0.03), (0.04, 0.07), (0.08, 0.06)]] S = [13, 26, 52] scaled_anchors = ( torch.tensor(anchors) * torch.tensor(S).unsqueeze(1).unsqueeze(1).repeat(1, 3, 2)) loss_expected = torch.tensor([1.0]) # WHEN loss = YoloLoss() loss_actual = loss(predictions=predictions, target=target, anchors=scaled_anchors[0]) # THEN self.assertEqual(loss_expected.item(), loss_actual.item())
def main(): model = YOLOv1(grid_size=GRID_SIZE, num_boxes=NUM_BOXES, num_classes=NUM_CLASSES).to(DEVICE) optimizer = optim.Adam( model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY ) loss_fn = YoloLoss(S=GRID_SIZE, B=NUM_BOXES, C=NUM_CLASSES) epochs_passed = 0 if LOAD_CHECKPOINT: checkpoint = torch.load(LOAD_CHECKPOINT_PATH) model.load_state_dict(checkpoint['model']) optimizer.load_state_dict(checkpoint['optimizer']) for param_group in optimizer.param_groups: param_group['lr'] = 2e-6 epochs_passed = checkpoint['epoch'] mask_dataset = YOLOVOCDataset( TRAINING_DATA, transform=train_transform, img_dir=IMG_DIR, label_dir=LABEL_DIR, S=GRID_SIZE, B=NUM_BOXES, C=NUM_CLASSES, hflip_prob = 0.5, random_crops=0.25, ) train_dataset, val_dataset = random_split ( mask_dataset, [TRAIN_SIZE, VAL_SIZE], generator=torch.Generator().manual_seed(42), ) if torch.cuda.is_available(): train_loader = DataLoader( dataset=train_dataset, batch_size=BATCH_SIZE, num_workers=NUM_WORKERS, pin_memory=PIN_MEMORY, shuffle=True, drop_last=DROP_LAST, ) val_loader = DataLoader( dataset=val_dataset, batch_size=BATCH_SIZE, num_workers=NUM_WORKERS, pin_memory=PIN_MEMORY, shuffle=True, drop_last=False, ) else: train_loader = DataLoader( dataset=train_dataset, batch_size=BATCH_SIZE, #num_workers=NUM_WORKERS, #pin_memory=PIN_MEMORY, shuffle=True, drop_last=DROP_LAST, ) val_loader = DataLoader( dataset=val_dataset, batch_size=BATCH_SIZE, #num_workers=NUM_WORKERS, #pin_memory=PIN_MEMORY, shuffle=True, drop_last=False, ) for epoch in range(EPOCHS): if epochs_passed%2 == 0: train_pred_boxes, train_target_boxes = get_bboxes( train_loader, model, iou_threshold=0.35, prob_threshold=0.4, S=GRID_SIZE, C=NUM_CLASSES, mode = "batch", device=DEVICE, ) val_pred_boxes, val_target_boxes = get_bboxes( val_loader, model, iou_threshold=0.35, prob_threshold=0.4, S=GRID_SIZE, C=NUM_CLASSES, mode = "batch", device=DEVICE, ) # map function takes predicted boxes and ground truth # boxes in form [[],[],[],...] where each sublist is a bounding box # of form [image_index, class_pred, x_mid, y_mid, w, h, prob] train_mAP = single_map( train_pred_boxes, train_target_boxes, box_format="midpoint", num_classes=NUM_CLASSES, iou_threshold=0.35, ) val_mAP = single_map( val_pred_boxes, val_target_boxes, box_format="midpoint", num_classes=NUM_CLASSES, iou_threshold=0.35, ) train_mAPs.append(train_mAP) val_mAPs.append(val_mAP) epochs_recorded.append(epochs_passed) print("Train mAP: %f"%(train_mAP)) print("Val mAP: %f"%(val_mAP)) if epochs_passed%10 == 0: save_path = SAVE_CHECKPOINT_PATH + ("checkpoint_%de"%(epochs_passed))+".pt" checkpoint = { 'epoch': epochs_passed, 'model': model.state_dict(), 'optimizer': optimizer.state_dict(), } torch.save(checkpoint,save_path) print("Trained for %d epochs"%(epochs_passed)) train_fn(train_loader, model, optimizer, loss_fn) epochs_passed += 1
def main(): """ Setup the model, loss function, data loader. Run the train function for each epoch. We may want to consider tune the number for non max suppression and mean average precision in order to remove false positive. BY false positive, they are a set of wrong output boxes. We can get rid of them by increasing the parameters. """ model = YOLOv3(num_classes=config.NUM_CLASSES).to(config.DEVICE) optimizer = optim.Adam(model.parameters(), lr=config.LEARNING_RATE, weight_decay=config.WEIGHT_DECAY) loss_fn = YoloLoss() scaler = torch.cuda.amp.GradScaler() train_loader, test_loader, train_eval_loader = get_loaders( train_csv_path=config.DATASET + "/train.csv", test_csv_path=config.DATASET + "/test.csv") if config.LOAD_MODEL: load_checkpoint(config.CHECKPOINT_FILE, model, optimizer, config.LEARNING_RATE) scaled_anchors = (torch.tensor(config.ANCHORS) * torch.tensor( config.S).unsqueeze(1).unsqueeze(1).repeat(1, 3, 2)).to(config.DEVICE) for epoch in range(config.NUM_EPOCHS): # plot_couple_examples(model, test_loader, 0.6, 0.5, scaled_anchors) train_fn(train_loader, model, optimizer, loss_fn, scaler, scaled_anchors) if config.SAVE_MODEL: save_checkpoint(model, optimizer, filename=f"checkpoint.pth.tar") # print(f"Currently epoch {epoch}") # print("On Train Eval loader:") # print("On Train loader:") # check_class_accuracy(model, train_loader, threshold=config.CONF_THRESHOLD) if epoch > 0 and epoch % 3 == 0: check_class_accuracy(model, test_loader, threshold=config.CONF_THRESHOLD) pred_boxes, true_boxes = get_evaluation_bboxes( test_loader, model, iou_threshold=config.NMS_IOU_THRESH, anchors=config.ANCHORS, threshold=config.CONF_THRESHOLD, ) mapval = mean_average_precision( pred_boxes, true_boxes, iou_threshold=config.MAP_IOU_THRESH, box_format="midpoint", num_classes=config.NUM_CLASSES, ) print(f"MAP: {mapval.item()}") model.train()
def train_loop(cfg_path, gpu_n='0'): # get configs with open(cfg_path, 'r') as stream: config = yaml.safe_load(stream) device = torch.device('cuda:{}'.format(gpu_n) if config['GPU'] and torch.cuda.is_available else 'cpu') dtype = torch.float32 # TODO: find out how it affects speed and accuracy MODEL = config['MODEL'] LOAD_MODEL = config['LOAD_MODEL'] LOAD_MODEL_FILE = config['LOAD_MODEL_FILE'] SAVE_MODEL = config['SAVE_MODEL'] SAVE_MODEL_N = config['SAVE_MODEL_N'] SAVE_MODEL_DIR = config['SAVE_MODEL_DIR'] DATASET_DIR = config['DATASET_DIR'] L_RATE = config['LEARNING_RATE'] DECAY_RATE = config['DECAY_RATE'] DECAY_EPOCHS = config['DECAY_EPOCHS'] WEIGHT_DECAY = config['WEIGHT_DECAY'] EPOCHS = config['EPOCHS'] BATCH_SIZE = config['BATCH_SIZE'] NUM_WORKERS = config['NUM_WORKERS'] PIN_MEMORY = config['PIN_MEMORY'] CSV_TRAIN = config['CSV_TRAIN'] CSV_VAL = config['CSV_VAL'] # set up model if MODEL == 'Darknet': model = YoloV1(grid_size=7, num_boxes=2, num_classes=20).to(DEVICE) elif MODEL == 'VGG': pass # add here VGG backbone if LOAD_MODEL: # TODO: load backbone # cfg_cp, start_epoch = load_checkpoint(LOAD_MODEL_FILE, model) val = input( 'Do you want to use config from checkpoint? Answer "yes" or "no": ' ) # if 'val' == 'yes': # L_RATE = cfg_cp['LEARNING_RATE'] # DECAY_RATE = cfg_cp['DECAY_RATE'] # DECAY_EPOCHS = cfg_cp['DECAY_EPOCHS'] # WEIGHT_DECAY = cfg_cp['WEIGHT_DECAY'] # BALANCED = cfg_cp['BALANCED_DATASET'] # BATCH_SIZE = cfg_cp['BATCH_SIZE'] # NUM_WORKERS = cfg_cp['NUM_WORKERS'] # PIN_MEMORY = cfg_cp['PIN_MEMORY'] # MIN_IMAGES = cfg_cp['MIN_IMAGES'] # LOSS = cfg_cp['LOSS'] else: model = init_weights(model) start_epoch = 0 optimizer = optim.Adam(model.parameters(), lr=L_RATE, weight_decay=WEIGHT_DECAY) loss_fn = YoloLoss() loader_params = BATCH_SIZE, NUM_WORKERS, PIN_MEMORY, DATASET_DIR, CSV_TRAIN, CSV_VAL loader = get_dataloader(loader_params) # create folder to save models if SAVE_MODEL: if not os.path.exists('{}/{}'.format(SAVE_MODEL_DIR, MODEL)): os.makedirs('{}/{}'.format(SAVE_MODEL_DIR, MODEL)) losses, accuracies = { 'train': [], 'validate': [] }, { 'train': [], 'validate': [] } for epoch in range(start_epoch, EPOCHS + start_epoch): t = time() if (epoch + 1) % DECAY_EPOCHS == 0: L_RATE *= (1 - DECAY_RATE) optimizer = optim.Adam(model.parameters(), lr=L_RATE, weight_decay=WEIGHT_DECAY) # print epoch number print_report(part='start', epoch=epoch) # train loop train_epoch(loader['train'], model, optimizer, device, loss_fn) # print metrics pred_bb, target_bb = get_bboxes(loader['train'], model, iou_threshold=0.5, threshold=0.4) train_map = mean_average_precision(pred_bb, target_bb, iou_threshold=0.5, box_format='midpoint') v_pred_bb, v_target_bb = get_bboxes(loader['val'], model, iou_threshold=0.5, threshold=0.4) val_map = mean_average_precision(v_pred_bb, v_target_bb, iou_threshold=0.5, box_format='midpoint') metrics = -1, -1, train_map, val_map print_report(part='accuracy', metrics=metrics) # collect metrics # losses['train'].append(train_loss) # losses['validate'].append(val_loss) # accuracies['train'].append(train_acc) # accuracies['validate'].append(val_acc) # save models # if SAVE_MODEL: # save_checkpoint(model=model, cfg=cfg, epoch=epoch, loss=round(val_loss, 3)) # print time print_report(part='end', t=int(time() - t))
def main(): model = Yolov1(split_size=S, num_boxes=B, num_classes=C).to(DEVICE) optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY) loss_fn = YoloLoss(S=S, B=B, C=C) if LOAD_MODEL: load_checkpoint(torch.load(LOAD_MODEL_FILE), model, optimizer) train_dataset = VOCDataset( training_path= '/home/mt/Desktop/For_github/computer_vision_projects/face_recognition/data', S=3, C=2, transform=transform) # test_dataset = VOCDataset( # "data/test.csv", transform=transform, img_dir=IMG_DIR, label_dir=LABEL_DIR, # ) train_loader = DataLoader( dataset=train_dataset, batch_size=BATCH_SIZE, num_workers=NUM_WORKERS, pin_memory=PIN_MEMORY, shuffle=True, drop_last=True, ) # test_loader = DataLoader( # dataset=test_dataset, # batch_size=BATCH_SIZE, # num_workers=NUM_WORKERS, # pin_memory=PIN_MEMORY, # shuffle=True, # drop_last=True, # ) for epoch in range(EPOCHS): # for x, y in train_loader: # x = x.to(DEVICE) # for idx in range(8): # bboxes = cellboxes_to_boxes(model(x)) # bboxes = non_max_suppression(bboxes[idx], iou_threshold=0.5, threshold=0.4, box_format="midpoint") # plot_image(x[idx].permute(1,2,0).to("cpu"), bboxes) # import sys # sys.exit() pred_boxes, target_boxes = get_bboxes(train_loader, model, iou_threshold=0.5, threshold=0.4) mean_avg_prec = mean_average_precision(pred_boxes, target_boxes, iou_threshold=0.5, box_format="midpoint") print(f"Train mAP: {mean_avg_prec}") #if mean_avg_prec > 0.9: # checkpoint = { # "state_dict": model.state_dict(), # "optimizer": optimizer.state_dict(), # } # save_checkpoint(checkpoint, filename=LOAD_MODEL_FILE) # import time # time.sleep(10) train_fn(train_loader, model, optimizer, loss_fn)
def train_loop(cfg_path, gpu_n='0', stat_path='stat'): # get configs with open(cfg_path, 'r') as stream: config = yaml.safe_load(stream) print() print(config) print() device = torch.device('cuda:{}'.format(gpu_n) if config['GPU'] and torch.cuda.is_available else 'cpu') dtype = torch.float32 # TODO: find out how it affects speed and accuracy MODEL = config['MODEL'] LOAD_MODEL = config['LOAD_MODEL'] LOAD_MODEL_FILE = config['LOAD_MODEL_FILE'] SAVE_MODEL = config['SAVE_MODEL'] SAVE_MODEL_N = config['SAVE_MODEL_N'] SAVE_MODEL_DIR = config['SAVE_MODEL_DIR'] DATASET_DIR = config['DATASET_DIR'] EPOCHS = config['EPOCHS'] BATCH_SIZE = config['BATCH_SIZE'] NUM_WORKERS = config['NUM_WORKERS'] PIN_MEMORY = config['PIN_MEMORY'] CSV_TRAIN = config['CSV_TRAIN'] CSV_VAL = config['CSV_VAL'] OPTIMIZER = config['OPTIMIZER'] # create stat file nid = create_stat(stat_path, config) # set up model S, B, C = 7, 2, 20 # TODO: add it to config if LOAD_MODEL: # TODO: load backbone model, cfg_save, epoch = load_checkpoint( LOAD_MODEL_FILE, device=device, S=S, B=B, C=C, cfg=config if MODEL == 'VGG16' else None) # TODO: init weight else: if MODEL == 'Darknet': model = YoloV1(grid_size=S, num_boxes=B, num_classes=C).to(device) elif MODEL == 'VGG': pass # add here VGG backbone model = init_weights(model) start_epoch = 0 loss_fn = YoloLoss() loader_params = BATCH_SIZE, NUM_WORKERS, PIN_MEMORY, DATASET_DIR, CSV_TRAIN, CSV_VAL loader = get_dataloader(loader_params, my_transforms) # create folder to save models if SAVE_MODEL: if not os.path.exists('{}/{}'.format(SAVE_MODEL_DIR, MODEL)): os.makedirs('{}/{}'.format(SAVE_MODEL_DIR, MODEL)) losses, accuracies = { 'train': [], 'validate': [] }, { 'train': [], 'validate': [] } optimizer = None opt_lr = None for epoch in range(start_epoch, EPOCHS + start_epoch): t = time() optimizer, opt_name, opt_lr = get_optimizer(optimizer, model, OPTIMIZER, epoch, opt_lr) print_report(part='start', epoch=epoch) # train loop train_epoch(loader['train'], model, optimizer, device, loss_fn, (opt_name, opt_lr)) # print metrics train_loss, train_maps = get_metrics_NEW(loader=loader['train'], model=model, iou_threshold=0.5, threshold=0.4, device=device, loss_func=loss_fn, S=S, B=B, C=C) val_loss, val_maps = get_metrics_NEW(loader=loader['val'], model=model, iou_threshold=0.5, threshold=0.4, device=device, loss_func=loss_fn) metrics = train_loss, val_loss, train_maps, val_maps print_report(part='accuracy', metrics=metrics) # collect metrics losses['train'].append(train_loss) losses['validate'].append(val_loss) accuracies['train'].append(np.mean(train_maps)) accuracies['validate'].append(np.mean(val_maps)) # write stats to CSV r = [ nid, datetime.now(), epoch, train_loss, val_loss, train_maps, val_maps ] with open('{}/stat.csv'.format(stat_path), 'a') as f: writer = csv.writer(f) writer.writerow(r) # save models # if SAVE_MODEL: # save_checkpoint(model=model, cfg=cfg, epoch=epoch, loss=round(val_loss, 3)) # print time print_report(part='end', t=int(time() - t))
def main_(): model = Yolov1(split_size=Grid, num_boxes=BBOX, num_classes=C_num).to(DEVICE) model_opt = Yolov1(split_size=Grid, num_boxes=BBOX, num_classes=C_num).to(DEVICE) best_loss = 1e6 # print(model) optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY) loss_fn = YoloLoss() train_dataset = YoloDataSet( "data/train.csv", transform=transform, img_dir=IMG_DIR, label_dir=LABEL_DIR, ) test_dataset = YoloDataSet( "data/test.csv", transform=transform, img_dir=IMG_DIR, label_dir=LABEL_DIR, ) train_loader = DataLoader( dataset=train_dataset, batch_size=BATCH_SIZE, num_workers=NUM_WORKERS, pin_memory=PIN_MEMORY, shuffle=True, drop_last=True, ) test_loader = DataLoader( dataset=test_dataset, batch_size=BATCH_SIZE, num_workers=NUM_WORKERS, pin_memory=PIN_MEMORY, shuffle=True, drop_last=True, ) test_class = [] test_conf = [] test_win = [] test_noob = [] test_cen = [] train_class = [] train_conf = [] train_win = [] train_noob = [] train_cen = [] for epoch in range(EPOCHS): print("epoch: " + str(epoch)) mean_loss = [] mean_coloss = [] mean_clloss = [] mean_celoss = [] mean_wloss = [] mean_noloss = [] model.eval() for batch_idx, (x, y) in enumerate(test_loader): x, y = x.to(DEVICE), y.to(DEVICE) with torch.no_grad(): out = model(x) # x, y = x_train.to(DEVICE), y_train.to(DEVICE) # out = model(x) loss, confidence_loss, class_loss, center_loss, window_loss, noobj_loss = loss_fn( out, y) mean_loss.append(loss.item()) mean_coloss.append(confidence_loss.item()) mean_clloss.append(class_loss.item()) mean_celoss.append(center_loss.item()) mean_wloss.append(window_loss.item()) mean_noloss.append(noobj_loss.item()) test_loss = sum(mean_loss) / len(mean_loss) if test_loss < best_loss: model_opt = model best_loss = test_loss model.train() print(f"test loss was {test_loss}") # print(f"Mean confidance loss was {sum(mean_coloss)/len(mean_coloss)}") # print(f"Mean class loss was {sum(mean_clloss)/len(mean_clloss)}") # print(f"Mean center loss was {sum(mean_celoss)/len(mean_celoss)}") # print(f"Mean window loss was {sum(mean_wloss)/len(mean_wloss)}") # print(f"Mean noobj loss was {sum(mean_noloss)/len(mean_noloss)}") test_class.append(sum(mean_clloss) / len(mean_clloss)) test_conf.append(sum(mean_coloss) / len(mean_coloss)) test_cen.append(sum(mean_celoss) / len(mean_celoss)) test_win.append(sum(mean_wloss) / len(mean_wloss)) test_noob.append(sum(mean_noloss) / len(mean_noloss)) # pred_boxes, target_boxes = get_bboxes(test_loader, model, iou_threshold=0.5, threshold=0.4, S=Grid, B=BBOX, C=C_num) # mean_avg_prec = mean_average_precision_built(pred_boxes, target_boxes, iou_threshold=0.5, box_format="midpoint") # print(f"test mAP: {mean_avg_prec}") # pred_boxes, target_boxes = get_bboxes(train_loader, model, iou_threshold=0.5, threshold=0.4, S=Grid, B=BBOX, C=C_num) # mean_avg_prec = mean_average_precision_built(pred_boxes, target_boxes, iou_threshold=0.5, box_format="midpoint") # print(f"Train mAP: {mean_avg_prec}") conf_loss, cla_loss, cen_loss, win_loss, noo_los = train_fn( train_loader, model, optimizer, loss_fn) train_class.append(cla_loss) train_conf.append(conf_loss) train_cen.append(cen_loss) train_win.append(win_loss) train_noob.append(noo_los) ploLoss(epoch, train_class, test_class, "class") ploLoss(epoch, train_conf, test_conf, "confidence") ploLoss(epoch, train_win, test_win, "win") ploLoss(epoch, train_noob, test_noob, "noobj") ploLoss(epoch, train_cen, test_cen, "center") model.eval() torch.save(model_opt, PATH)
save_checkpoint, load_checkpoint, ) from loss import YoloLoss import numpy as np LEARNING_RATE = 2e-5 WEIGHT_DECAY = 0 DEVICE = "cuda" if torch.cuda.is_available else "cpu" model = Yolov1(split_size=7, num_boxes=2, num_classes=20).to(DEVICE) optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY) loss_fn = YoloLoss() IMG_DIR = "data/images" LABEL_DIR = "data/labels" NUM_WORKERS = 0 BATCH_SIZE = 1 PIN_MEMORY = True class Compose(object): def __init__(self, transforms): self.transforms = transforms def __call__(self, img, bboxes): for t in self.transforms: img, bboxes = t(img), bboxes
def main(): # Data loading train_csv_path = config.DATASET + "100examples.csv" test_csv_path = config.DATASET + "100examples.csv" train_dataset = YoloDataset(train_csv_path, transforms=config.max_transforms, Scale=config.Scale, image_dir=config.IMG_DIR, label_dir=config.LABEL_DIR, anchor_boxes=config.ANCHORS, number_of_anchors=3, number_of_scales=3, ignore_iou_threshold=0.5, num_anchors_per_scale=3) test_dataset = YoloDataset(test_csv_path, transforms=config.max_transforms, Scale=config.Scale, image_dir=config.IMG_DIR, label_dir=config.LABEL_DIR, anchor_boxes=config.ANCHORS, number_of_anchors=3, number_of_scales=3, ignore_iou_threshold=0.5, num_anchors_per_scale=3) train_loader = DataLoader(dataset=train_dataset, batch_size=config.BATCH_SIZE, num_workers=config.NUM_WORKERS, pin_memory=config.PIN_MEMORY, shuffle=False, drop_last=True) test_loader = DataLoader(dataset=test_dataset, batch_size=config.BATCH_SIZE, num_workers=config.NUM_WORKERS, pin_memory=config.PIN_MEMORY, shuffle=False, drop_last=True) # Model model = YoloV3(num_classes=config.C).to(device=config.DEVICE) optimizer = optim.Adam(model.parameters(), lr=config.LEARNING_RATE, weight_decay=config.WEIGHT_DECAY) loss_function = YoloLoss() scalar = torch.cuda.amp.GradScaler() # Miscellaneous scaled_anchors = (torch.tensor(config.ANCHORS) * torch.tensor( config.Scale).unsqueeze(1).unsqueeze(1).repeat(1, 3, 2)).to( config.DEVICE) writer = SummaryWriter() current_time = time.time() # Loading previously saved model weights if config.LOAD_MODEL: load_checkpoint(config.TRAINED_FILE, model, optimizer, config.LEARNING_RATE) # Epoch loop for epoch in range(config.NUM_EPOCHS): print("Epoch:", epoch) # Run training model.train() losses = [] for x, y in train_loader: x = x.float() x = x.to(config.DEVICE) y0, y1, y2 = (y[0].to(config.DEVICE), y[1].to(config.DEVICE), y[2].to(config.DEVICE)) with torch.cuda.amp.autocast(): y_p = model(x) loss_0 = loss_function(predictions=y_p[0], target=y0, anchors=scaled_anchors[0]) loss_1 = loss_function(predictions=y_p[1], target=y1, anchors=scaled_anchors[1]) loss_2 = loss_function(predictions=y_p[2], target=y2, anchors=scaled_anchors[2]) loss = loss_0 + loss_1 + loss_2 losses.append(loss.item()) optimizer.zero_grad() scalar.scale(loss).backward() scalar.step(optimizer) scalar.update() mean_loss = sum(losses) / len(losses) writer.add_scalar("loss: ", mean_loss, epoch) # Run evaluation if config.SAVE_MODEL and epoch > 0 and epoch % config.SAVE_FREQUENCY == 0: batch_recalls, batch_precisions = [], [] model.eval() for x, y in test_loader: x = x.float() x = x.to(config.DEVICE) with torch.no_grad(): yp = model(x) # Move predictions to cpu yp = [yp[0].to('cpu'), yp[1].to('cpu'), yp[2].to('cpu')] # boxes_from_yp(yp) returns all yp bboxes in a batch yp_boxes = boxes_from_yp(yp=yp, iou_threshold=config.MAP_IOU_THRESH, threshold=config.CONF_THRESHOLD) # boxes_from_y(y) returns all y bboxes in a batch y_boxes = boxes_from_y(y=y) # Calculate precision for batch precision, recall = calc_batch_precision_recall( y_boxes=y_boxes, yp_boxes=yp_boxes, iou_threshold=config.MAP_IOU_THRESH) batch_precisions.append(precision) batch_recalls.append(recall) mean_precision = sum(batch_precisions) / len(batch_precisions) mean_recall = sum(batch_recalls) / len(batch_recalls) # Log evaluation variables writer.add_scalar("precision", mean_precision, epoch) writer.add_scalar("recall", mean_recall, epoch) if config.RENDER and epoch > 0 and epoch % config.SAVE_FREQUENCY == 0: # Render prediction bboxes x, y = next(x for i, x in enumerate(test_loader) if i == 5) model.eval() with torch.no_grad(): x_gpu = x.float() x_gpu = x_gpu.to(config.DEVICE) yp = model(x_gpu) yp = [yp[0].to('cpu'), yp[1].to('cpu'), yp[2].to('cpu')] x = draw_yp_on_x(x, yp, probability_threshold=0.5, anchors=config.ANCHORS) x = draw_y_on_x(x, y) grid = torchvision.utils.make_grid(x, nrow=4) writer.add_image("yp on x", grid, global_step=epoch) # Save model if config.SAVE_MODEL and epoch > 0 and epoch % config.SAVE_FREQUENCY == 0: save_checkpoint(model, optimizer, filename=config.CHECKPOINT_FILE) delta_time, current_time = time_function(current_time) writer.add_scalar("Epoch Duration [s]", delta_time, epoch) writer.flush()
def main(): model = Yolov1(split_size=7, num_boxes=2, num_classes=20).to(DEVICE) optimizer = optim.Adam( model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY ) loss_fn = YoloLoss() if LOAD_MODEL: load_checkpoint(torch.load(LOAD_MODEL_FILE), model, optimizer) train_dataset = VOCDataset( # test.csv, 8examples.csv, 100examples.csv "data/8examples.csv", transform=transform, img_dir=IMG_DIR, label_dir=LABEL_DIR, ) test_dataset = VOCDataset( "data/test.csv" , transform = transform, img_dir = IMG_DIR, label_dir = LABEL_DIR ) train_loader = DataLoader( dataset=train_dataset, batch_size=BATCH_SIZE, num_workers=NUM_WORKERS, pin_memory=PIN_MEMORY, shuffle=True, drop_last=False, ) test_loader = DataLoader( dataset=test_dataset, batch_size=BATCH_SIZE, num_workers=NUM_WORKERS, pin_memory=PIN_MEMORY, shuffle=True, drop_last=False, ) for epoch in range(EPOCHS): # for x, y in train_loader: # x = x.to(DEVICE) # for idx in range(8): # bboxes = cellboxes_to_boxes(model(x)) # bboxes = non_max_suppression(bboxes[idx], iou_threshold=0.5, threshold=0.4, box_format="midpoint") # plot_image(x[idx].permute(1,2,0).to("cpu"), bboxes) # import sys # sys.exit() pred_boxes, target_boxes = get_bboxes( train_loader, model, iou_threshold=0.5, threshold=0.4 , device= DEVICE, ) mean_avg_prec = mean_average_precision( pred_boxes, target_boxes, iou_threshold=0.5, box_format="midpoint" ) if mean_avg_prec > 0.9: checkpoint = { "state_dict": model.state_dict(), "optimizer": optimizer.state_dict(), } save_checkpoint(checkpoint, filename=LOAD_MODEL_FILE) import time time.sleep(10) print(f"Train mAP: {mean_avg_prec}") train_fn(train_loader, model, optimizer, loss_fn)
import os os._exit(1) # Enable anomaly detection for debugging purpose torch.autograd.set_detect_anomaly(True) model = YOLO((channels, height, width), S, B, C) if num_gpu > 1: print("Let's use", num_gpu, "GPUs!") model = torch.nn.DataParallel(model) model = model.to(device) # Init tensorboard for loss visualization tb = Tensorboard() loss_func = YoloLoss(S, B, C) # torch.nn.MSELoss(reduction="sum") optimizer = opt(model.parameters(), momentum=momentum, lr=lr, weight_decay=weight_decay) # momentum=momentum # Keep the loss of the best model best_model_loss = test = float("inf") start_time = time.ctime() print("### start_time:", start_time) cur_epoch = 0 for epoch in range(cur_epoch, epochs): if epoch == 74 or epoch == 104: print("Changing learning from {} to {}...".format(lr, lr * 0.1))
def main(): IMG_DIR = "../datasets/voc/images" LABEL_DIR = "../datasets/voc/labels" MAPPING_FILE = "../datasets/voc/100examples.csv" BATCH_SIZE = 16 NUM_WORKERS = 2 PIN_MEMORY = True DEVICE = "cuda" if torch.cuda.is_available() else "cpu" LEARNING_RATE = 2e-4 WEIGHT_DECAY = 0.0 EPOCHS = 100 LOAD_MODEL_FILE = "retrained_grad_true.pth.tar" S = 7 C = 20 B = 2 PRETRAINED = True print(DEVICE) class Compose(object): def __init__(self, transforms): self.transforms = transforms def __call__(self, image, bboxes): for t in self.transforms: image = t(image) return image, bboxes # Types of preprocessing transforms we want to apply convert_transform = transforms.ToTensor() resize_transform = transforms.Resize((448, 448)) # normalize_transform = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) transform = Compose([convert_transform, resize_transform]) train_dataset = VOCDataset(MAPPING_FILE, S=S, C=C, B=B, transform=transform, img_dir=IMG_DIR, label_dir=LABEL_DIR) train_loader = DataLoader(dataset=train_dataset, batch_size=BATCH_SIZE, num_workers=NUM_WORKERS, pin_memory=PIN_MEMORY, shuffle=True, drop_last=True) val_dataset = VOCDataset(MAPPING_FILE, S=S, C=C, B=B, transform=transform, img_dir=IMG_DIR, label_dir=LABEL_DIR) val_loader = DataLoader(dataset=val_dataset, batch_size=BATCH_SIZE, num_workers=NUM_WORKERS, pin_memory=PIN_MEMORY, shuffle=True, drop_last=True) model = YoloV1(S=S, B=B, C=C, pretrained=PRETRAINED).to(DEVICE) for child in model.resnet18.parameters(): child.requires_grad = True model.resnet18.fc.requires_grad = True optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY) yolo_loss = YoloLoss(S=S, C=C, B=B) print("=> Loading checkpoint") model.load_state_dict(torch.load(LOAD_MODEL_FILE)["state_dict"]) optimizer.load_state_dict(torch.load(LOAD_MODEL_FILE)["optimizer"]) model.train() # Sets model to training mode model.eval() # Sets model to evaluation mode with torch.no_grad(): # Use with evaluation mode for epoch in tqdm(range(1)): for i, (x, y) in enumerate(val_loader): x, y = x.to(DEVICE), y.to(DEVICE) y_p = model(x) x = x.to('cpu') y_p = y_p.to('cpu') y_p = torch.reshape(y_p, (BATCH_SIZE, S, S, C + 5 * B)) print("Rendering original labels (y to x)") render_batch(x, y, S=S, B=B, C=C) print("Rendering predicted labels (y_p to x)") render_batch(x, y_p, S=S, B=B, C=C)
def main(): # Data loading test_csv_path = config.DATASET + "100examples.csv" test_dataset = YoloDataset(test_csv_path, transforms=config.max_transforms, Scale=config.Scale, image_dir=config.IMG_DIR, label_dir=config.LABEL_DIR, anchor_boxes=config.ANCHORS, number_of_anchors=3, number_of_scales=3, ignore_iou_threshold=0.5, num_anchors_per_scale=3) test_loader = DataLoader(dataset=test_dataset, batch_size=config.BATCH_SIZE, num_workers=config.NUM_WORKERS, pin_memory=config.PIN_MEMORY, shuffle=False, drop_last=True) # Model model = YoloV3(num_classes=config.C) optimizer = optim.Adam(model.parameters(), config.LEARNING_RATE, weight_decay=config.WEIGHT_DECAY) loss_function = YoloLoss() scalar = torch.cuda.amp.GradScaler() writer = SummaryWriter() model.to(device=config.DEVICE) model.eval() def load_checkpoint(checkpoint_file, model, optimizer, lr): print("=> Loading checkpoint") checkpoint = torch.load(checkpoint_file, map_location=config.DEVICE) model.load_state_dict(checkpoint["state_dict"]) optimizer.load_state_dict(checkpoint["optimizer"]) # If we don't do this then it will just have learning rate of old checkpoint # and it will lead to many hours of debugging \: for param_group in optimizer.param_groups: param_group["lr"] = lr load_checkpoint(config.CHECKPOINT_FILE, model, optimizer, config.LEARNING_RATE) # all_pred_boxes, all_true_boxes = get_evaluation_bboxes(test_loader, model, iou_threshold=0.5, anchors=config.ANCHORS, threshold=0.5, box_format="midpoint", device="cuda") model.eval() train_idx = 0 anchors = config.ANCHORS threshold = 0.9 from utils_depreciated import cells_to_bboxes from utils_depreciated import non_max_suppression for batch_idx, (x, labels) in enumerate(test_loader): x = x.float() x = x.to("cuda") all_pred_boxes = [] all_true_boxes = [] with torch.no_grad(): predictions = model(x) batch_size = x.shape[0] bboxes = [[] for _ in range(batch_size)] for i in range(3): S = predictions[i].shape[2] anchor = torch.tensor([*anchors[i]]).to("cuda") * S boxes_scale_i = cells_to_bboxes(predictions[i], anchor, S=S, is_preds=True) for idx, (box) in enumerate(boxes_scale_i): bboxes[idx] += box # we just want one bbox for each label, not one for each scale true_bboxes = cells_to_bboxes(labels[2], anchor, S=S, is_preds=False) for idx in range(batch_size): nms_boxes = non_max_suppression(bboxes[idx], iou_threshold=0.9, threshold=0.9, box_format="midpoint") for nms_box in nms_boxes: all_pred_boxes.append([train_idx] + nms_box) for box in true_bboxes[idx]: if box[1] > threshold: all_true_boxes.append([train_idx] + box) train_idx += 1 from torchvision_utils import draw_bounding_boxes x = x.to('cpu') * 255 #boxes = all_pred_boxes #cells_to_boxes_render(y) boxes = [(torch.tensor(x)[3:] * torch.tensor([416, 416, 416, 416])).tolist() for x in all_pred_boxes] boxes = torch.tensor(boxes) for idx in range(x.shape[0]): x[idx, ...] = draw_bounding_boxes(image=x[idx, ...].type(torch.uint8), boxes=boxes, colors=None, width=2, pad_value=2) grid = torchvision.utils.make_grid(x) writer.add_image("images", grid, global_step=batch_idx) writer.flush() """
def main(): # Data loading train_dataset = CustomDataset(root=config.root_train, annFile=config.annFile_train, transforms=config.train_transforms, catagory=config.CATEGORY_FILTER) val_dataset = CustomDataset(root=config.root_val, annFile=config.annFile_val, transforms=config.val_transforms, catagory=config.CATEGORY_FILTER) train_loader = DataLoader(dataset=train_dataset, batch_size=8, num_workers=2, pin_memory=True, shuffle=True, drop_last=True) val_loader = DataLoader(dataset=val_dataset, batch_size=8, num_workers=2, pin_memory=True, shuffle=False, drop_last=True) # Model model = YoloV3(num_classes=config.C).to(device=config.DEVICE) # from model_external import YOLOv3 # model = YOLOv3(num_classes=90).to(device=config.DEVICE) optimizer = optim.Adam(model.parameters(), lr=config.LEARNING_RATE, weight_decay=config.WEIGHT_DECAY) from loss_external import YoloLoss loss_function = YoloLoss().to(device=config.DEVICE) # Miscellaneous scaled_anchors = (torch.tensor(config.anchors) * torch.tensor( config.Scale).unsqueeze(1).unsqueeze(1).repeat(1, 3, 2)).to( config.DEVICE) writer = SummaryWriter() current_time = time.time() # Loading previously saved model weights if config.LOAD_MODEL: load_checkpoint("cp.pth.tar", model, optimizer, config.LEARNING_RATE) print(torch.cuda.memory_summary(device=None, abbreviated=False)) # Training loop for cycle in range(config.CYCLES): print("Cycle:", cycle) x, y = next(iter(val_loader)) x = x.to(config.DEVICE) y0, y1, y2 = (y[0].to(config.DEVICE), y[1].to(config.DEVICE), y[2].to(config.DEVICE)) yp = model(x) loss_0 = loss_function(predictions=yp[0], target=y0, anchors=scaled_anchors[0]) loss_1 = loss_function(predictions=yp[1], target=y1, anchors=scaled_anchors[1]) loss_2 = loss_function(predictions=yp[2], target=y2, anchors=scaled_anchors[2]) loss = loss_0 + loss_1 + loss_2 loss.backward() optimizer.step() optimizer.zero_grad() # Run validation if cycle % 100 == 0 and cycle != 0: model.eval() losses = [] with torch.no_grad(): x, y = next(iter(val_loader)) x = x.to(config.DEVICE) y0, y1, y2 = (y[0].to(config.DEVICE), y[1].to(config.DEVICE), y[2].to(config.DEVICE)) yp = model(x) loss_0 = loss_function(predictions=yp[0], target=y0, anchors=scaled_anchors[0]) loss_1 = loss_function(predictions=yp[1], target=y1, anchors=scaled_anchors[1]) loss_2 = loss_function(predictions=yp[2], target=y2, anchors=scaled_anchors[2]) loss = loss_0 + loss_1 + loss_2 losses.append(loss) avg_val_loss = sum(losses) / len(losses) writer.add_scalar("val_loss: ", avg_val_loss, cycle) model.train() # Run validation """ if cycle % 100 == 0 and cycle != 0: model.eval() x, y = next(iter(val_loader)) x = x.float() x = x.to(config.DEVICE) # y0, y1, y2 = (y[0].to(config.DEVICE), y[1].to(config.DEVICE), y[2].to(config.DEVICE)) with torch.no_grad(): yp = model(x) # Move predictions to cpu yp = [yp[0].to('cpu'), yp[1].to('cpu'), yp[2].to('cpu')] # boxes_from_yp(yp) returns all yp bboxes in a batch yp_boxes = boxes_from_yp(yp=yp, iou_threshold=config.MAP_IOU_THRESH, threshold=config.CONF_THRESHOLD) # boxes_from_y(y) returns all y bboxes in a batch y_boxes = boxes_from_y(y=y) """ # Save model if cycle % 1000 == 0 and cycle != 0: save_checkpoint(model, optimizer, cycle, filename=config.CHECKPOINT_FILE) # Rendering loop if cycle % 100 == 0 and cycle != 0: model.eval() x, y = next(iter(val_loader)) with torch.no_grad(): x_gpu = x.to(config.DEVICE) yp = model(x_gpu) yp = [yp[0].to('cpu'), yp[1].to('cpu'), yp[2].to('cpu')] x = denormalize(x) * 255 draw_y_on_x(x, y) draw_yp_on_x(x, yp, probability_threshold=0.5, anchors=config.anchors) # Save batch grid as image image_dir = "./batch_dir" image_dir_exists = os.path.exists(image_dir) if not image_dir_exists: os.makedirs(image_dir) img_name = str(image_dir) + "/batch_" + str(cycle) + ".png" save_image(x / 255, img_name) model.train() writer.add_scalar("train_loss: ", loss.item(), cycle) delta_time, current_time = time_function(current_time) writer.add_scalar("Epoch Duration [s]", delta_time, cycle) writer.flush()
def main(): IMG_DIR = "../datasets/voc/images" LABEL_DIR = "../datasets/voc/labels" MAPPING_FILE = "../datasets/voc/100examples.csv" BATCH_SIZE = 8 NUM_WORKERS = 2 PIN_MEMORY = True DEVICE = "cuda" if torch.cuda.is_available() else "cpu" LEARNING_RATE = 4e-4 WEIGHT_DECAY = 0.0 EPOCHS = 100 LOAD_MODEL_FILE = "retrained_grad_true.pth.tar" S = 7 C = 20 B = 2 PRETRAINED = True print("Using Device:", DEVICE) class Compose(object): def __init__(self, transforms): self.transforms = transforms def __call__(self, image, bboxes): for t in self.transforms: image = t(image) return image, bboxes # Types of preprocessing transforms we want to apply convert_transform = transforms.ToTensor() resize_transform = transforms.Resize((448, 448)) # normalize_transform = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) transform = Compose([convert_transform, resize_transform]) train_dataset = VOCDataset(MAPPING_FILE, S=S, C=C, B=B, transform=transform, img_dir=IMG_DIR, label_dir=LABEL_DIR) train_loader = DataLoader(dataset=train_dataset, batch_size=BATCH_SIZE, num_workers=NUM_WORKERS, pin_memory=PIN_MEMORY, shuffle=True, drop_last=True) val_dataset = VOCDataset(MAPPING_FILE, S=S, C=C, B=B, transform=transform, img_dir=IMG_DIR, label_dir=LABEL_DIR) val_loader = DataLoader(dataset=val_dataset, batch_size=BATCH_SIZE, num_workers=NUM_WORKERS, pin_memory=PIN_MEMORY, shuffle=True, drop_last=True) model = YoloV1(S=S, B=B, C=C, pretrained=PRETRAINED).to(DEVICE) for counter, param in enumerate(model.resnet.parameters()): param.requires_grad = True model.resnet.fc.requires_grad = True for counter, param in enumerate(model.parameters()): print(counter, param.requires_grad) parameters = [param.numel() for param in model.parameters()] print("Number of model parameters", sum(parameters)) parameters = [param.numel() for param in model.resnet.parameters()] print("Number of resnet parameters", sum(parameters)) parameter_list = filter(lambda p: p.requires_grad, model.parameters()) optimizer = optim.Adam(parameter_list, lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY) yolo_loss = YoloLoss(S=S, C=C, B=B) writer = SummaryWriter() current_time = time.time() mean_loss_list = [] for epoch in range(EPOCHS): mean_loss = [] pred_boxes, target_boxes = get_bboxes(train_loader, model, iou_threshold=0.5, threshold=0.4, device=DEVICE) mean_avg_prec = mean_average_precision(pred_boxes, target_boxes, iou_threshold=0.5, box_format="midpoint") # Save the model if mean_avg_prec > 0.99 or epoch == (EPOCHS - 1): checkpoint = { "state_dict": model.state_dict(), "optimizer": optimizer.state_dict(), } print("=> Saving checkpoint") torch.save(checkpoint, LOAD_MODEL_FILE) time.sleep(10) for batch_idx, (x, y) in enumerate(train_loader): x, y = x.to(DEVICE), y.to(DEVICE) y_p = model(x) # y = torch.flatten(y, start_dim=1, end_dim=3) y_p = torch.reshape(y_p, (BATCH_SIZE, S, S, C + 5 * B)) loss = yolo_loss(y_p, y) mean_loss.append(loss.item()) optimizer.zero_grad() loss.backward() optimizer.step() # Calculate average loss mean_loss = sum(mean_loss)/len(mean_loss) delta_time = time.time() - current_time current_time = time.time() writer.add_scalar("Average Loss: ", mean_loss, epoch) writer.add_scalar("Mean Average Precision: ", mean_avg_prec, epoch) writer.add_scalar("Epoch Duration [s]", delta_time, epoch) print("Epoch:", epoch) mean_loss_list.append(mean_loss) writer.close()
def criterion(self, y, d): loss = YoloLoss(20, self.anchors) return loss.forward(y, d)
def main(): TRAIN_IMG_DIR = "../datasets/coco/images/train2017/" VAL_IMG_DIR = "../datasets/coco/images/val2017/" TEST_IMG_DIR = "../datasets/coco/images/test2017/" LABEL_DIR = "../datasets/bbox_labels_train.json" BATCH_SIZE = 16 NUM_WORKERS = 2 PIN_MEMORY = True DEVICE = "cuda" if torch.cuda.is_available() else "cpu" LEARNING_RATE = 2e-4 WEIGHT_DECAY = 0.0 EPOCHS = 50 TRAIN = False LOAD_MODEL_FILE = "epoch.pth.tar" S = 7 C = 100 B = 2 print(DEVICE) class Compose(object): def __init__(self, transforms): self.transforms = transforms def __call__(self, image, bboxes): for t in self.transforms: image = t(image) return image, bboxes # Types of preprocessing transforms we want to apply convert_transform = transforms.ToTensor() resize_transform = transforms.Resize((448, 448)) transform = Compose([convert_transform, resize_transform]) train_dataset = COCODataset("../datasets/bbox_labels_val.json", S=S, C=C, B=B, transform=transform, img_dir=VAL_IMG_DIR, label_dir=LABEL_DIR) train_loader = DataLoader(dataset=train_dataset, batch_size=BATCH_SIZE, num_workers=NUM_WORKERS, pin_memory=PIN_MEMORY, shuffle=True, drop_last=True) val_dataset = COCODataset("../datasets/bbox_labels_val.json", S=S, C=C, B=B, transform=transform, img_dir=VAL_IMG_DIR, label_dir=LABEL_DIR) val_loader = DataLoader(dataset=val_dataset, batch_size=BATCH_SIZE, num_workers=NUM_WORKERS, pin_memory=PIN_MEMORY, shuffle=True, drop_last=True) model = YoloV1(S=S, B=B, C=C).to(DEVICE) optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY) yolo_loss = YoloLoss(S=S, C=C, B=B) writer = SummaryWriter() # TRAINING if TRAIN: mean_loss_list = [] for epoch in range(EPOCHS): mean_loss = [] for batch_idx, (x, y) in enumerate(train_loader): x, y = x.to(DEVICE), y.to(DEVICE) y_p = model(x) # y = torch.flatten(y, start_dim=1, end_dim=3) y_p = torch.reshape(y_p, (BATCH_SIZE, S, S, C + 5 * B)) loss = yolo_loss(y_p, y) mean_loss.append(loss.item()) optimizer.zero_grad() loss.backward() optimizer.step() # Calculate average loss mean_loss = sum(mean_loss) / len(mean_loss) writer.add_scalar("Average Loss: ", mean_loss, epoch) print("Epoch:", epoch) print(f"Mean loss was {mean_loss}") mean_loss_list.append(mean_loss) # Save the model if True: checkpoint = { "state_dict": model.state_dict(), "optimizer": optimizer.state_dict(), } print("=> Saving checkpoint") torch.save(checkpoint, LOAD_MODEL_FILE) time.sleep(10) writer.close() plt.plot(mean_loss_list) plt.show() # VALIDATION else: print("=> Loading checkpoint") model.load_state_dict(torch.load(LOAD_MODEL_FILE)["state_dict"]) optimizer.load_state_dict(torch.load(LOAD_MODEL_FILE)["optimizer"]) model.train() # Sets model to training mode model.eval() # Sets model to evaluation mode with torch.no_grad(): # Use with evaluation mode for epoch in tqdm(range(1)): for i, (x, y) in enumerate(val_loader): x, y = x.to(DEVICE), y.to(DEVICE) y_p = model(x) x = x.to('cpu') y_p = y_p.to('cpu') y_p = torch.reshape(y_p, (BATCH_SIZE, S, S, C + 5 * B)) print("Rendering original labels (y to x)") render_batch(x, y) print("Rendering predicted labels (y_p to x)") render_batch(x, y_p)