batch_size=batch_size,
shuffle=True,
drop_last=True,
num_workers=8)
test_dset = VOCDetection(root=data_root, split='test')
test_dloader = DataLoader(test_dset,
batch_size=batch_size,
shuffle=False,
drop_last=False,
num_workers=8)
model = Yolo(grid_size, num_boxes, num_classes)
model = model.to(device)
pretrained_weights = torch.load(pretrained_backbone_path)
model.load_state_dict(pretrained_weights)
print('loaded pretrained weight')
# Freeze the backbone network.
model.features.requires_grad_(False)
model_params = [v for v in model.parameters() if v.requires_grad is True]
optimizer = optim.SGD(model_params, lr=lr, momentum=0.9, weight_decay=5e-4)
compute_loss = Loss(grid_size, num_boxes, num_classes)
# Load the last checkpoint if exits.
ckpt_path = os.path.join(ckpt_dir, 'last.pth')
if os.path.exists(ckpt_path):
ckpt = torch.load(ckpt_path)
model.load_state_dict(ckpt['model'])
optimizer.load_state_dict(ckpt['optimizer'])
args = parser.parse_args()
print(args)
device = torch.device(
'cuda' if args.use_gpu and torch.cuda.is_available() else 'cpu')
if not os.path.exists(args.output_path):
os.makedirs(args.output_path)
# Initiate model
model = Yolo(num_classes=20).to(device)
# If specified we start from checkpoint
if args.pretrained_weights:
if args.pretrained_weights.endswith('.pth'):
model.load_state_dict(torch.load(args.pretrained_weights))
else:
model.load_darknet_weights(args.pretrained_weights)
# Get dataloader
train_dataset = VOCDetection(args.train_path, args.img_size)
train_loader = DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
collate_fn=train_dataset.collate_fn)
val_dataset = VOCDetection(args.val_path, args.img_size)
val_loader = DataLoader(val_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.num_workers,