def main(): #check device use_cuda = torch.cuda.is_available() torch.manual_seed(123) device = torch.device('cuda:1' if use_cuda else 'cpu') print('Device used:', device) #model create and to device model = models.Yolov1_vgg16bn(pretrained=True).to(device) #dataset init and loader warper #windows trainset = localLoader.AERIAL('hw2_train_val/train15000', transform=transforms.ToTensor()) trainset_loader = DataLoader(trainset, batch_size=32, shuffle=True, num_workers=1, collate_fn=collate_fn) testset = localLoader.AERIAL('hw2_train_val/val1500', transform=transforms.ToTensor()) testset_loader = DataLoader(testset, batch_size=32, shuffle=False, num_workers=1, collate_fn=collate_fn) #Debug(model, testset_loader) train(model, trainset_loader, testset_loader, device, 1000)
def main(): if not os.path.exists(args.output): os.mkdir(args.output) torch.set_default_dtype(torch.float) device = utils.selectDevice() # Initialize model if args.command == "basic": model = models.Yolov1_vgg16bn(pretrained=True).to(device) elif args.command == "improve": model = models.Yolov1_vgg16bn_Improve(pretrained=True).to(device) model = utils.loadModel(args.model, model) model.eval() # Initialize dataset transform = transforms.Compose([ transforms.Resize((448, 448)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) # TODO: Reset as MyDataset testset = dataset.Testset(img_root=args.images, transform=transform) test_loader = DataLoader(testset, batch_size=1, shuffle=False, num_workers=4) # Iterative interference for batch_idx, (data, imgName) in enumerate(test_loader, 1): data = data.to(device) output = model(data) if args.command == "basic": boxes, classIndexs, probs = decode(output, nms=args.nms, prob_min=args.prob, iou_threshold=args.iou) if args.command == "improve": boxes, classIndexs, probs = decode(output, nms=args.nms, prob_min=args.prob, iou_threshold=args.iou, grid_num=14) classNames = labelEncoder.inverse_transform(classIndexs.type(torch.long).to("cpu")) export(boxes, classNames, probs, imgName[0] + ".txt", args.output) print("Predicted: [{}/{} ({:.2%})]\r".format( batch_idx, len(test_loader.dataset), batch_idx / len(test_loader.dataset)), end="" )
test_data = Load_Data(test_label, test_fname, test_img, test_iname, train=False) # In[16]: # trainset_loader = DataLoader(train_data, batch_size=32, shuffle=False, num_workers=1) testset_loader = DataLoader(test_data, batch_size=1, shuffle=False, num_workers=1) checkpoint_path = "/home/robot/hw2-Chee-An-Yu/pass0.092/yolo-2500.pth" model = models.Yolov1_vgg16bn(pretrained=True) model.eval() optimizer = optim.SGD(model.parameters(), lr=0.00001, momentum=0.9) device = torch.device('cuda') model.to(device) print('load model...') load_checkpoint(checkpoint_path, model, optimizer) # model.load_state_dict(torch.load('yolo-5500.pth')) #trainset = Yolotrainset('hw2_train_val/train15000/') #trainloader = DataLoader(trainset, batch_size=1, shuffle=False) # valset = Yolovalset('hw2_train_val/val1500/') # valloader = DataLoader(valset, batch_size=1, shuffle=False) print('load valset end') # model.eval()
def main(): transform = transforms.Compose([ transforms.Resize((448, 448)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) grid_num = 7 if args.command == "basic" else 14 trainset = dataset.MyDataset(root="hw2_train_val/train15000", grid_num=grid_num, train=args.augment, transform=transform) testset = dataset.MyDataset(grid_num=grid_num, root="hw2_train_val/val1500", train=False, transform=transform) trainLoader = DataLoader(trainset, batch_size=args.batchs, shuffle=True, num_workers=args.worker) testLoader = DataLoader(testset, batch_size=1, shuffle=False, num_workers=args.worker) device = utils.selectDevice(show=True) if args.command == "basic": model = models.Yolov1_vgg16bn(pretrained=True).to(device) criterion = models.YoloLoss(7., 2., 5., 0.5, device).to(device) optimizer = optim.SGD(model.parameters(), lr=args.lr, weight_decay=1e-4) scheduler = optim.lr_scheduler.MultiStepLR(optimizer, [20, 45, 55], gamma=0.1) start_epoch = 0 if args.load: model, optimizer, start_epoch, scheduler = utils.loadCheckpoint( args.load, model, optimizer, scheduler) model = train(model, criterion, optimizer, scheduler, trainLoader, testLoader, start_epoch, args.epochs, device, lr=args.lr, grid_num=7) elif args.command == "improve": model_improve = models.Yolov1_vgg16bn_Improve( pretrained=True).to(device) criterion = models.YoloLoss(14., 2., 5, 0.5, device).to(device) optimizer = optim.SGD(model_improve.parameters(), lr=args.lr, weight_decay=1e-4) scheduler = optim.lr_scheduler.MultiStepLR(optimizer, [20, 40, 70], gamma=0.1) start_epoch = 0 if args.load: model_improve, optimizer, start_epoch, scheduler = utils.loadCheckpoint( args.load, model, optimizer, scheduler) model_improve = train(model_improve, criterion, optimizer, scheduler, trainLoader, testLoader, start_epoch, args.epochs, device, lr=args.lr, grid_num=7, save_name="Yolov1-Improve")
action='store', dest='test_dir', default='./hw2_train_val/val1500/images/') parser.add_argument('-p', action='store', dest='pred_dir', default='./hw2_train_val/val1500/predTxt_hbb/') parser.add_argument('-g', action='store', dest='grid_size', default=7) parser.add_argument('-ep', action='store', dest='epochs', default=30) args = parser.parse_args() if int(args.grid_size) == 8: print('Running improved model') model = models.Improved_model(pretrained=True).to(device) else: model = models.Yolov1_vgg16bn(pretrained=True).to(device) if args.train: print('Training start') train(model, int(args.epochs), grid_size=int(args.grid_size)) else: if args.ckpt_path == None: print("Need ckpt path!") elif args.eval: print('Test with: ', args.ckpt_path.split('/')[-1]) optimizer = optim.Adam(model.parameters(), lr=0.0001) state = torch.load(args.ckpt_path) model.load_state_dict(state['state_dict']) optimizer.load_state_dict(state['optimizer']) gen_pred_file(model, pred_path=args.pred_dir,