def load(self, fname=None):
if fname is None:
fnames = [
f.split('_')[-1]
for f in glob(os.path.join(self.__checkpoints_dir, "*.pth"))
]
fname = self.__model_name + '_' + max(fnames)
elif fname.split('.')[-1] != ".pth":
fname = fname + ".pth"
checkpoint_path = os.path.join(cfg.get_checkpoints_root(), fname)
if cfg.get_device().type == "cpu":
self.load_state_dict(
torch.load(checkpoint_path, map_location="cpu"))
else:
self.load_state_dict(torch.load(checkpoint_path))
if __name__ == "__main__":
update_cfg()
filter_dataset()
dataset = BEVKitti()
loader = DataLoader(dataset=dataset,
batch_size=cfg.get_batch_size(),
shuffle=True,
num_workers=cfg.get_num_workers(),
pin_memory=cfg.get_pin_memory())
model = ComplexYOLO().to(cfg.get_device())
optimizer = optim.SGD(model.parameters(),
lr=cfg.get_initial_lr(),
momentum=cfg.get_momentum(),
weight_decay=cfg.get_weight_decay())
region_loss = RegionLoss()
dataset_size = len(dataset)
for epoch in range(cfg.get_n_epochs()):
for group in optimizer.param_groups:
if epoch >= 4 and epoch < 80:
group["lr"] = 1e-4
elif epoch >= 80 and epoch < 160:
group["lr"] = 1e-5
elif epoch >= 160:
group["lr"] = 1e-6
plt.imshow(bev)
# Draw target
target = target[target[:, 0] != 0]
for obj in target:
draw_box(obj[1:7], bev_H, bev_W, color='r')
# Draw prediction
for obj in preds:
draw_box(obj[0:6], bev_H, bev_W, color='g')
plt.show()
if __name__ == "__main__":
update_cfg()
device = cfg.get_device()
val_dataset = BEVKitti("val")
model = ComplexYOLO().to(device)
model.load()
model.eval()
for i, (bev, target) in enumerate(val_dataset):
bev = bev.unsqueeze(0)
bev = bev.to(device)
output = model(bev)
bev = bev.squeeze(0)
bev = bev.cpu().detach().numpy()
target = target.cpu().detach().numpy()
prediction_bev = get_prediction_bev(output, bev.shape[1, 2])
prediction_bev = prediction_bev.cpu().detach().numpy()
def __init__(self):
super(RegionLoss, self).__init__()
device = cfg.get_device()
self.__mse_loss = nn.MSELoss(reduction="mean").to(device)
self.__bce_loss = nn.BCELoss(reduction="mean").to(device)
self.__ce_loss = nn.CrossEntropyLoss().to(device)