def train():
# step.1
# model = getattr(models, config.model)().cuda()
model = getattr(models, config.model)()
model = torch.nn.DataParallel(model) # multi-gpu
model.cuda()
print 'train', config.model
# print model
if config.load_model_path:
model.load_state_dict(torch.load(config.load_model_path))
# step.2
train_data = KITTIRoadFusion(config.root, split='train', num_features=19)
val_data = KITTIRoadFusion(config.root, split='val', num_features=19)
train_dataloader = DataLoader(train_data,
batch_size=config.batch_size,
shuffle=True,
num_workers=config.num_workers,
drop_last=True)
val_dataloader = DataLoader(val_data, num_workers=config.num_workers)
# step.3
# binary cross entropy 2d
criterion = BooststrapBinaryCrossEntropyLoss2D(
int(config.IMG_SIZE * config.K1))
if config.optim is 'Adam':
optimizer = optim.Adam(model.parameters(),
weight_decay=config.weight_decay) # optimizer
else:
optimizer = optim.SGD(model.parameters(),
lr=config.learning_rate,
momentum=0.9,
weight_decay=config.weight_decay)
best_max_f = 100
# step.4 train
for epoch in range(config.max_epoch):
for i, (_, _, cloud, _, _, lb) in enumerate(train_dataloader):
cloud, lb = Variable(cloud), Variable(lb)
cloud, lb = cloud.float().cuda(), lb.long().cuda()
optimizer.zero_grad()
# forward
pred = model(cloud)
loss = criterion(pred, lb)
# backward
loss.backward()
optimizer.step()
if i % config.print_freq == 0:
print 'epoch: %03d,iter: %03d, loss=%lf' % (epoch, i,
loss.data[0])
if (epoch + 1) % config.test_freq == 0:
max_f = val(model, val_dataloader)
if max_f > best_max_f:
best_max_f = max_f
print 'better model, MaxF = %f' % best_max_f
# save the best model
filename = '%sK_%.2f.pth' % (config.model_prefix, config.K1)
torch.save(model.state_dict(), filename)
print 'save %s succeed!' % filename
return best_max_f
def train():
# step.1
# model = getattr(models, config.model)().cuda()
model = getattr(models, config.model)()
model = torch.nn.DataParallel(model) # multi-gpu
model.cuda()
print 'train', config.model
# weights initialize
# model.load_state_dict(weights_initialize(model.state_dict()))
# print model
if config.load_model_path:
model.load_state_dict(torch.load(config.load_model_path))
# step.2 数据集
train_data = KITTIRoadFusion(config.root, split='train', num_features=19)
val_data = KITTIRoadFusion(config.root, split='val', num_features=19)
train_dataloader = DataLoader(train_data, batch_size=config.batch_size, shuffle=True, num_workers=config.num_workers, drop_last=True)
val_dataloader = DataLoader(val_data, num_workers=config.num_workers)
# step.3
# binary cross entropy 2d
criterion_coarse = BooststrapBinaryCrossEntropyLoss2D(int(config.IMG_SIZE/16*config.K2))
criterion_fine = BooststrapBinaryCrossEntropyLoss2D(int(config.IMG_SIZE*config.K1))
if config.optim is 'Adam':
optimizer = optim.Adam(model.parameters(), weight_decay=config.weight_decay) # optimizer
else:
optimizer = optim.SGD(model.parameters(), lr=config.learning_rate, momentum=0.9, weight_decay=config.weight_decay)
best_max_f = 0
# step.4 train
for epoch in range(config.max_epoch):
for i, (_, im, cloud, theta, shift, lb) in enumerate(train_dataloader):
im, cloud, theta, shift, lb = Variable(im), Variable(cloud), Variable(theta), Variable(shift), Variable(lb)
im, cloud, theta, shift, lb = im.float().cuda(), cloud.float().cuda(), theta.float().cuda(), shift.float().cuda(), lb.long().cuda()
optimizer.zero_grad()
# forward
pred_coarse, pred_fine = model(im, cloud, theta, shift)
loss_coarse = criterion_coarse(pred_coarse, lb[:, ::4, ::4])
loss_fine = criterion_fine(pred_fine, lb)
loss = config.LAMBDA*loss_coarse + (1-config.LAMBDA)*loss_fine
# backward
loss.backward()
optimizer.step()
if i % config.print_freq == 0:
print 'epoch: %03d,iter: %03d, loss_coarse=%lf, loss_fine=%lf, loss=%lf' \
% (epoch, i, loss_coarse.data[0], loss_fine.data[0], loss.data[0])
if (epoch+1) % config.test_freq == 0:
max_f = val(model, val_dataloader)
if max_f > best_max_f:
best_max_f = max_f
print 'better model, MaxF = %f' % best_max_f
# save the best model
filename = '%sK1_%.2f-K2_%.2f-LAMBDA_%.2f.pth' % (config.model_prefix, config.K1, config.K2, config.LAMBDA)
torch.save(model.state_dict(), filename)
print 'save %s succeed!' % filename
return best_max_f
def eval_on_validation():
if not os.path.exists(config.run_dir):
os.mkdir(config.run_dir)
model = getattr(models, config.model)()
model = torch.nn.DataParallel(model) # multi-gpu
model.cuda()
print 'test', config.model
print model
if config.load_model_path:
model.load_state_dict(torch.load(config.load_model_path))
# data
test_data = KITTIRoadFusion(config.root, split='val', num_features=19)
test_dataloader = DataLoader(test_data,
batch_size=1,
shuffle=False,
num_workers=4)
# test
model.eval()
total_time = 0
for i, (name, _, cloud, _, _, _) in enumerate(test_dataloader):
cloud = Variable(cloud)
cloud = cloud.float().cuda()
start = time.clock()
pred = model(cloud) # inference
end = time.clock()
total_time += (end - start)
pred = pred.data.cpu().numpy().squeeze()
pred = np.array(pred * 255, dtype=np.uint8)
filename = os.path.join(config.run_dir, name[0])
print filename
# cv2.imwrite(filename, pred)
print 'Runtime ############# time(s) : %f ##########' % (
total_time / test_dataloader.__len__())
def eval_on_validation():
if not os.path.exists(config.run_dir):
os.mkdir(config.run_dir)
model = getattr(models, config.model)()
model = torch.nn.DataParallel(model) # multi-gpu
model.cuda()
print 'test on validation set.', config.model
print model
if config.load_model_path:
model.load_state_dict(torch.load(config.load_model_path))
# data
test_data = KITTIRoadFusion(config.root, split='val', num_features=19)
test_dataloader = DataLoader(test_data, batch_size=1, shuffle=False, num_workers=4)
# test
model.eval()
eval_acc = 0
eval_acc_cls = 0
eval_mean_iu = 0
eval_fwavacc = 0
li_pred = []
li_gt = []
total_time = 0
for i, (name, im, cloud, theta, shift, lb) in enumerate(test_dataloader):
im, cloud, theta, shift, lb = Variable(im), Variable(cloud), Variable(theta), Variable(shift), Variable(lb)
im, cloud, theta, shift, lb = im.float().cuda(), cloud.float().cuda(), theta.float().cuda(), shift.float().cuda(), lb.long().cuda()
start = time.clock()
_, pred = model(im, cloud, theta, shift) # inference
end = time.clock()
total_time += (end-start)
# pred = F.upsample_bilinear(pred, scale_factor=4)
# save image
label_pred = pred.data.cpu().numpy().squeeze()
label_pred = np.array(label_pred*255, dtype=np.uint8)
filename = os.path.join(config.run_dir, name[0])
print filename
# cv2.imwrite(filename, label_pred)
# Mean IoU
label_true = lb.data.cpu().numpy().astype(np.int8)
label_pred = pred.data.cpu().numpy().squeeze(0)
label_pred = (label_pred > 0.5).astype(np.int8)
for (label, prob) in zip(label_true, label_pred):
acc, acc_cls, mean_iu, fwavacc = label_accuracy_score(label, prob, n_class=2)
eval_acc += acc
eval_acc_cls += acc_cls
eval_mean_iu += mean_iu
eval_fwavacc += fwavacc
# MaxF
label_pred = pred.data.cpu().numpy().squeeze()
label_true = lb.data.cpu().numpy().squeeze()
li_pred.append(label_pred)
li_gt.append(label_true)
print 'Runtime ############# time(s) : %f ##########' % (total_time / test_dataloader.__len__())
print 'Validation ======ACC: %lf,Mean IoU: %lf======' % (eval_acc / test_dataloader.__len__(),
eval_mean_iu / test_dataloader.__len__())
eval_road(li_pred, li_gt)
def eval_on_validation_bev():
if not os.path.exists(config.run_dir):
os.mkdir(config.run_dir)
model = getattr(models, config.model)()
model = torch.nn.DataParallel(model) # multi-gpu
model.cuda()
print 'test on validation set.', config.model
print model
if config.load_model_path:
model.load_state_dict(torch.load(config.load_model_path))
# data
bev = BirdsEyeView()
test_data = KITTIRoadFusion(config.root, split='val', num_features=19, return_bev=True)
test_dataloader = DataLoader(test_data, batch_size=1, shuffle=False, num_workers=4)
# test
model.eval()
eval_acc = 0
eval_acc_cls = 0
eval_mean_iu = 0
eval_fwavacc = 0
li_pred = []
li_gt = []
for i, (name, im, cloud, theta, shift, _, lb) in enumerate(test_dataloader):
im, cloud, theta, shift, lb = Variable(im), Variable(cloud), Variable(theta), Variable(shift), Variable(lb)
im, cloud, theta, shift, lb = im.float().cuda(), cloud.float().cuda(), theta.float().cuda(), shift.float().cuda(), lb.long().cuda()
_, pred = model(im, cloud, theta, shift) # inference
pred = pred.data.cpu().numpy().squeeze()
theta = theta.data.cpu().numpy().squeeze()
shift = shift.data.cpu().numpy().squeeze()
label_pred = bev.transformLable2BEV((pred*255).astype(np.uint8), theta, shift)
label_true = lb.data.cpu().numpy().squeeze()
# save image
filename = os.path.join(config.run_dir, name[0])
print filename
cv2.imwrite(filename, label_pred)
label_pred = label_pred/255.
# Mean IoU
label_true_hard = np.expand_dims(label_true.astype(np.int8), axis=0)
label_pred_hard = np.expand_dims((label_pred > 0.5).astype(np.int8), axis=0)
for (label, prob) in zip(label_true_hard, label_pred_hard):
acc, acc_cls, mean_iu, fwavacc = label_accuracy_score(label, prob, n_class=2)
eval_acc += acc
eval_acc_cls += acc_cls
eval_mean_iu += mean_iu
eval_fwavacc += fwavacc
# MaxF
li_pred.append(label_pred)
li_gt.append(label_true)
print 'Validation ======ACC: %lf,Mean IoU: %lf======' % (eval_acc / test_dataloader.__len__(),
eval_mean_iu / test_dataloader.__len__())
eval_road(li_pred, li_gt)
def eval():
if not os.path.exists(config.run_dir):
os.mkdir(config.run_dir)
model = getattr(models, config.model)()
model = torch.nn.DataParallel(model) # multi-gpu
model.cuda()
print 'test', config.model
print model
if config.load_model_path:
model.load_state_dict(torch.load(config.load_model_path))
# data
test_data = KITTIRoadFusion(config.root, split='test', num_features=19)
test_dataloader = DataLoader(test_data, batch_size=1, shuffle=False, num_workers=4)
# test
model.eval()
for i, (name, im, cloud, theta, shift) in enumerate(test_dataloader):
im, cloud, theta, shift = Variable(im), Variable(cloud), Variable(theta), Variable(shift)
im, cloud, theta, shift = im.float().cuda(), cloud.float().cuda(), theta.float().cuda(), shift.float().cuda()
_, pred = model(im, cloud, theta, shift) # inference
pred = pred.data.cpu().numpy().squeeze()
pred = np.array(pred*255, dtype=np.uint8)
filename = os.path.join(config.run_dir, name[0])
print filename
cv2.imwrite(filename, pred)