cfg.width = cfg.size cfg.height = cfg.size cfg.w = cfg.size cfg.h = cfg.size print(cfg) os.environ["CUDA_VISIBLE_DEVICES"] = cfg.gpu device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') logging.info(f'Using device {device}') if cfg.use_darknet_cfg: model = Darknet(cfg.cfgfile) else: model = Yolov4(cfg.pretrained, n_classes=cfg.classes) if cfg.load is not None: model.load_model(cfg.load, device) if cfg.pretrained is not None: model.load_model(cfg.pretrained, device, pretrained=True) #model.load_weights(cfg.pretrained) if torch.cuda.device_count() > 1: model = torch.nn.DataParallel(model) model.to(device=device) try: train(model=model, config=cfg, epochs=cfg.TRAIN_EPOCHS, device=device, log_step=20)
if dataset == 'patch': from cfg_patch import Cfg checkpoint = 'checkpoints/Yolov4_modanet_128_epoch46.pth' elif dataset == 'modanet_whole': pass cfg = Cfg cfg.gpu = '0' os.environ["CUDA_VISIBLE_DEVICES"] = cfg.gpu device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') json_gt = os.path.expanduser('~/data/datasets/modanet/Annots/modanet_instances_val_new.json') if cfg.use_darknet_cfg: model = Darknet(cfg.cfgfile) else: model = Yolov4(cfg.pretrained, n_classes=cfg.classes) model.load_model(checkpoint, device) if torch.cuda.device_count() > 1: model = torch.nn.DataParallel(model) model.to(device=device) #model.load_state_dict(torch.load(checkpoint), strict=True) val_dataset = YoloModanetHumanDataset(cfg.anno_path, cfg, train=False) n_val = len(val_dataset) val_loader = DataLoader(val_dataset, batch_size=1, shuffle=True, num_workers=0, pin_memory=True, drop_last=False, collate_fn=val_collate) evaluate_nms(model, val_loader, cfg, device, human_patch=True, json_gt=None) #evaluate_nms(model, val_loader, cfg, device, human_patch=True, json_gt=json_gt)
cfg.width = cfg.size cfg.height = cfg.size cfg.w = cfg.size cfg.h = cfg.size print(cfg) os.environ["CUDA_VISIBLE_DEVICES"] = cfg.gpu device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') logging.info(f'Using device {device}') if cfg.use_darknet_cfg: model = Darknet(cfg.cfgfile) else: model = Yolov4(cfg.pretrained, n_classes=cfg.classes) if cfg.load is not None: model.load_model(cfg.load, device) if torch.cuda.device_count() > 1: model = torch.nn.DataParallel(model) model.to(device=device) try: train(model=model, config=cfg, epochs=cfg.TRAIN_EPOCHS, device=device, ) except KeyboardInterrupt: torch.save(model.state_dict(), 'INTERRUPTED.pth') logging.info('Saved interrupt') try: sys.exit(0)