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 = 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():
model = Yolo(num_classes=config.NUM_CLASSES).to(config.DEVICE)
optimizer = optim.Adam(model.parameters(),
lr=config.LEARNING_RATE,
weight_decay=config.WEIGHT_DECAY)
lossfunction = 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.GRID_SIZE).unsqueeze(1).unsqueeze(1).repeat(1, 3, 2)).to(
config.DEVICE)
for epoch in range(config.NUM_EPOCHS):
train(train_loader, model, optimizer, lossfunction, scaler,
scaled_anchors)
if config.SAVE_MODEL:
save_checkpoint(model, optimizer, filename=config.CHECKPOINT_FILE)
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 main(args):
# scale normalized anchors by their corresponding scales (e.g img_size = 416, feature maps will 32, 16, 8,
# scaled_anchors = normalized_anchors X (8, 16, 32))
"""
scaled anchors
tensor([[ 3.6250, 2.8125],
[ 4.8750, 6.1875],
[11.6562, 10.1875],
[ 1.8750, 3.8125],
[ 3.8750, 2.8125],
[ 3.6875, 7.4375],
[ 1.2500, 1.6250],
[ 2.0000, 3.7500],
[ 4.1250, 2.8750]])
"""
di = torch.tensor([
int(args.img_size / 32),
int(args.img_size / 16),
int(args.img_size / 8)
]).unsqueeze(1)
scaled_anchors = (torch.tensor(anchors) * torch.repeat_interleave(
di, torch.tensor([3, 3, 3]), dim=0).repeat(1, 2)).to(args.device)
train_dataset = YOLODataset(args.train_img_dir,
args.train_label_dir,
args.train_annotation_file,
scaled_anchors,
num_classes,
args.iou_thre,
args.img_size,
transform=train_transforms)
val_dataset = YOLODataset(args.val_img_dir,
args.val_label_dir,
args.val_annotation_file,
scaled_anchors,
num_classes,
args.iou_thre,
args.img_size,
transform=test_transforms)
train_loader = DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True)
val_loader = DataLoader(val_dataset,
batch_size=args.batch_size,
shuffle=True)
model = YOLO(num_classes).to(args.device)
# if it is first time to train the model, use darknet53 pretrained weights
# load_weights_darknet53("darknet53.conv.74", model)
optimizer = optim.Adam(model.parameters(), lr=1e-5, weight_decay=1e-4)
load_checkpoint("checkpoint.pth.tar", model, optimizer, 1e-5)
for epoch in range(args.epoches):
train_loop(train_loader, model, YOLOloss, optimizer, args.device,
scaled_anchors)
print(f"Epoch {epoch+1}\n-------------------------------")
save_checkpoint(model, optimizer, filename=f"checkpoint.pth.tar")
model.eval()
results = check_class_accuracy(model, val_loader, args.conf_thre)
mAP = mean_average_precision(model, val_loader, scaled_anchors,
num_classes, args.iou_thre)
results.append(mAP)
with open("result.txt", 'a') as f:
f.write(str(results).strip("[").strip("]") + '\n')
model.train()
print("Done!")
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()
