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,