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)
