def get_model(cfg):
cfg.merge_from_file('../configs/second/car.yaml')
anchors = AnchorGenerator(cfg).anchors
preprocessor = Preprocessor(cfg)
model = Second(cfg).cuda().eval()
ckpt = torch.load('../pvrcnn/ckpts/epoch_12.pth')['state_dict']
model.load_state_dict(ckpt, strict=True)
return model, preprocessor, anchors
def main():
"""TODO: Trainer class to manage objects."""
model = Second(cfg).cuda()
parameters = model.parameters()
loss_fn = ProposalLoss(cfg)
preprocessor = TrainPreprocessor(cfg)
dataloader = build_train_dataloader(cfg, preprocessor)
optimizer = torch.optim.Adam(parameters, lr=0.01)
start_epoch = load_ckpt('./ckpts/epoch_10.pth', model, optimizer)
scheduler = build_lr_scheduler(optimizer, cfg, start_epoch,
len(dataloader))
train_model(model, dataloader, optimizer, scheduler, loss_fn,
cfg.TRAIN.EPOCHS, start_epoch)