def main():
global args, best_mIoU
args = parser.parse_args()
os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(str(gpu) for gpu in args.gpus)
args.gpus = len(args.gpus)
if args.no_partialbn:
sync_bn.Synchronize.init(args.gpus)
if args.dataset == 'VOCAug' or args.dataset == 'VOC2012' or args.dataset == 'COCO':
num_class = 21
ignore_label = 255
scale_series = [10, 20, 30, 60]
elif args.dataset == 'cityscapes':
num_class = 19
ignore_label = 0
scale_series = [15, 30, 45, 90]
else:
raise ValueError('Unknown dataset ' + args.dataset)
model = models.FCN(num_class, base_model=args.arch, dropout=args.dropout, partial_bn=not args.no_partialbn)
input_mean = model.input_mean
input_std = model.input_std
policies = model.get_optim_policies()
model = torch.nn.DataParallel(model, device_ids=range(args.gpus)).cuda()
if args.resume:
if os.path.isfile(args.resume):
print(("=> loading checkpoint '{}'".format(args.resume)))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
best_mIoU = checkpoint['best_mIoU']
torch.nn.Module.load_state_dict(model, checkpoint['state_dict'])
print(("=> loaded checkpoint '{}' (epoch {})".format(args.evaluate, checkpoint['epoch'])))
else:
print(("=> no checkpoint found at '{}'".format(args.resume)))
if args.weight:
if os.path.isfile(args.weight):
print(("=> loading initial weight '{}'".format(args.weight)))
checkpoint = torch.load(args.weight)
torch.nn.Module.load_state_dict(model, checkpoint['state_dict'])
else:
print(("=> no model file found at '{}'".format(args.weight)))
cudnn.benchmark = True
cudnn.fastest = True
# Data loading code
train_loader = torch.utils.data.DataLoader(
getattr(ds, args.dataset + 'DataSet')(data_list=args.train_list, transform=torchvision.transforms.Compose([
tf.GroupRandomHorizontalFlip(),
tf.GroupRandomScale(size=(0.5, 2.0), interpolation=(cv2.INTER_LINEAR, cv2.INTER_NEAREST)),
tf.GroupRandomCrop(size=args.train_size),
tf.GroupRandomPad(size=args.train_size, padding=(input_mean, (ignore_label, ))),
tf.GroupRandomRotation(degree=(-10, 10), interpolation=(cv2.INTER_LINEAR, cv2.INTER_NEAREST), padding=(input_mean, (ignore_label, ))),
tf.GroupRandomBlur(applied=(True, False)),
tf.GroupNormalize(mean=(input_mean, (0, )), std=(input_std, (1, ))),
])), batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True, drop_last=True)
val_loader = torch.utils.data.DataLoader(
getattr(ds, args.dataset + 'DataSet')(data_list=args.val_list, transform=torchvision.transforms.Compose([
tf.GroupCenterCrop(size=args.test_size),
tf.GroupConcerPad(size=args.test_size, padding=(input_mean, (ignore_label, ))),
tf.GroupNormalize(mean=(input_mean, (0, )), std=(input_std, (1, ))),
])), batch_size=args.batch_size * 3, shuffle=False, num_workers=args.workers, pin_memory=True)
# define loss function (criterion) optimizer and evaluator
criterion = torch.nn.NLLLoss(ignore_index=ignore_label).cuda()
for group in policies:
print(('group: {} has {} params, lr_mult: {}, decay_mult: {}'.format(group['name'], len(group['params']), group['lr_mult'], group['decay_mult'])))
optimizer = torch.optim.SGD(policies, args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
evaluator = EvalSegmentation(num_class, ignore_label)
if args.evaluate:
validate(val_loader, model, criterion, 0, evaluator)
return
for epoch in range(args.start_epoch, args.epochs):
adjust_learning_rate(optimizer, epoch, args.lr_steps)
# train for one epoch
train(train_loader, model, criterion, optimizer, epoch)
# evaluate on validation set
if (epoch + 1) % args.eval_freq == 0 or epoch == args.epochs - 1:
mIoU = validate(val_loader, model, criterion, (epoch + 1) * len(train_loader), evaluator)
# remember best mIoU and save checkpoint
is_best = mIoU > best_mIoU
best_mIoU = max(mIoU, best_mIoU)
save_checkpoint({
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'best_mIoU': best_mIoU,
}, is_best)
def main():
global args, best_mIoU
args = parser.parse_args()
os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(
str(gpu) for gpu in args.gpus)
args.gpus = len(args.gpus)
if args.dataset == 'VOCAug' or args.dataset == 'VOC2012' or args.dataset == 'COCO':
num_class = 21
ignore_label = 255
scale_series = [10, 20, 30, 60]
elif args.dataset == 'Cityscapes':
num_class = 19
ignore_label = 255
scale_series = [15, 30, 45, 90]
elif args.dataset == 'ApolloScape':
num_class = 37
ignore_label = 255
elif args.dataset == 'CULane' or args.dataset == 'L4E':
num_class = 5
ignore_label = 255
else:
raise ValueError('Unknown dataset ' + args.dataset)
model = models.ERFNet(3, num_class)
input_mean = model.input_mean
input_std = model.input_std
policies = model.get_optim_policies()
model = torch.nn.DataParallel(model, device_ids=range(args.gpus)).cuda()
if args.resume:
if os.path.isfile(args.resume):
print(("=> loading checkpoint '{}'".format(args.resume)))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
best_mIoU = checkpoint['best_mIoU']
torch.nn.Module.load_state_dict(model, checkpoint['state_dict'])
print(("=> loaded checkpoint '{}' (epoch {})".format(
args.evaluate, checkpoint['epoch'])))
else:
print(("=> no checkpoint found at '{}'".format(args.resume)))
cudnn.benchmark = True
cudnn.fastest = True
# Data loading code
test_loader = torch.utils.data.DataLoader(getattr(ds, 'VOCAugDataSet')(
data_list=args.val_list,
transform=torchvision.transforms.Compose([
tf.GroupRandomScaleNew(size=(args.img_width, args.img_height),
interpolation=(cv2.INTER_LINEAR,
cv2.INTER_NEAREST)),
tf.GroupNormalize(mean=(input_mean, (0, )),
std=(input_std, (1, ))),
])),
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=False)
# define loss function (criterion) optimizer and evaluator
weights = [1.0 for _ in range(5)]
weights[0] = 0.4
class_weights = torch.FloatTensor(weights).cuda()
criterion = torch.nn.NLLLoss(ignore_index=ignore_label,
weight=class_weights).cuda()
for group in policies:
print(('group: {} has {} params, lr_mult: {}, decay_mult: {}'.format(
group['name'], len(group['params']), group['lr_mult'],
group['decay_mult'])))
evaluator = EvalSegmentation(num_class, ignore_label)
### evaluate ###
validate(test_loader, model, criterion, 0, evaluator)
return
[255, 0, 0],
[0, 255, 0],
[0, 0, 255],
[255, 255, 0],
],
dtype=np.float32)
no_of_classes = 5
model = ENet_model(model_id,
img_height=img_height,
img_width=img_width,
batch_size=batch_size,
no_classes=no_of_classes)
train_mean_channels = pickle.load(open("data/mean_channels.pkl", 'rb'))
input_mean = train_mean_channels #[103.939, 116.779, 123.68] # [0, 0, 0]
input_std = [1, 1, 1]
normalizer = transforms.GroupNormalize(mean=(input_mean, (0, )),
std=(input_std, (1, )))
# load the mean color channels of the train imgs:
train_mean_channels = pickle.load(open("data/mean_channels.pkl", 'rb'))
# create a saver for restoring variables/parameters:
saver = tf.train.Saver(tf.trainable_variables(),
write_version=tf.train.SaverDef.V2)
sess = tf.Session()
# restore the best trained model:
saver.restore(sess, "./training_logs/model_augmented_best_epoch_ 12.ckpt")
logits = tf.nn.softmax(model.net.logits, -1)
exist_logits = model.net.line_existance_logit
exist_logits = tf.where(tf.less(exist_logits, 0.5), exist_logits * 0.0,
def main():
global args, best_mIoU
args = parser.parse_args()
os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(
str(gpu) for gpu in args.gpus)
#args.gpus = len(args.gpus)
if args.dataset == 'VOCAug' or args.dataset == 'VOC2012' or args.dataset == 'COCO':
num_class = 21
ignore_label = 255
scale_series = [10, 20, 30, 60]
elif args.dataset == 'Cityscapes':
num_class = 19
ignore_label = 255 # 0
scale_series = [15, 30, 45, 90]
elif args.dataset == 'ApolloScape':
num_class = 37 # merge the noise and ignore labels
ignore_label = 255
elif args.dataset == 'CULane':
num_class = 5
ignore_label = 255
else:
raise ValueError('Unknown dataset ' + args.dataset)
model = models.ERFNet(num_class)
input_mean = model.input_mean
input_std = model.input_std
model = model.cuda()
model = torch.nn.DataParallel(model, device_ids=args.gpus)
#model = torch.nn.DataParallel(model, device_ids=range(args.gpus)).cuda()
def load_my_state_dict(
model, state_dict
): # custom function to load model when not all dict elements
own_state = model.state_dict()
ckpt_name = []
cnt = 0
for name, param in state_dict.items():
if name not in list(own_state.keys()) or 'output_conv' in name:
ckpt_name.append(name)
continue
own_state[name].copy_(param)
cnt += 1
print('#reused param: {}'.format(cnt))
return model
if args.resume:
if os.path.isfile(args.resume):
print(("=> loading checkpoint '{}'".format(args.resume)))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
model = load_my_state_dict(model, checkpoint['state_dict'])
# torch.nn.Module.load_state_dict(model, checkpoint['state_dict'])
print(("=> loaded checkpoint '{}' (epoch {})".format(
args.evaluate, checkpoint['epoch'])))
else:
print(("=> no checkpoint found at '{}'".format(args.resume)))
cudnn.benchmark = True
cudnn.fastest = True
# Data loading code
train_loader = torch.utils.data.DataLoader(getattr(
ds,
args.dataset.replace("CULane", "VOCAug") + 'DataSet')(
data_list=args.train_list,
transform=torchvision.transforms.Compose([
tf.GroupRandomScale(size=(0.595, 0.621),
interpolation=(cv2.INTER_LINEAR,
cv2.INTER_NEAREST)),
tf.GroupRandomCropRatio(size=(args.img_width,
args.img_height)),
tf.GroupRandomRotation(degree=(-1, 1),
interpolation=(cv2.INTER_LINEAR,
cv2.INTER_NEAREST),
padding=(input_mean, (ignore_label, ))),
tf.GroupNormalize(mean=(input_mean, (0, )),
std=(input_std, (1, ))),
])),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=False,
drop_last=True)
val_loader = torch.utils.data.DataLoader(getattr(
ds,
args.dataset.replace("CULane", "VOCAug") + 'DataSet')(
data_list=args.val_list,
transform=torchvision.transforms.Compose([
tf.GroupRandomScale(size=(0.595, 0.621),
interpolation=(cv2.INTER_LINEAR,
cv2.INTER_NEAREST)),
tf.GroupRandomCropRatio(size=(args.img_width,
args.img_height)),
tf.GroupNormalize(mean=(input_mean, (0, )),
std=(input_std, (1, ))),
])),
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=False)
# define loss function (criterion) optimizer and evaluator
weights = [1.0 for _ in range(5)]
weights[0] = 0.4
class_weights = torch.FloatTensor(weights).cuda()
criterion = torch.nn.NLLLoss(ignore_index=ignore_label,
weight=class_weights).cuda()
criterion_exist = torch.nn.BCEWithLogitsLoss().cuda()
optimizer = torch.optim.SGD(model.parameters(),
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
evaluator = EvalSegmentation(num_class, ignore_label)
args.evaluate = False #True
if args.evaluate:
validate(val_loader, model, criterion, 0, evaluator)
return
for epoch in range(args.epochs): # args.start_epoch
adjust_learning_rate(optimizer, epoch, args.lr_steps)
# train for one epoch
train(train_loader, model, criterion, criterion_exist, optimizer,
epoch)
# evaluate on validation set
if (epoch + 1) % args.eval_freq == 0 or epoch == args.epochs - 1:
mIoU = validate(val_loader, model, criterion,
(epoch + 1) * len(train_loader), evaluator)
# remember best mIoU and save checkpoint
is_best = mIoU > best_mIoU
best_mIoU = max(mIoU, best_mIoU)
save_checkpoint(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'best_mIoU': best_mIoU,
}, is_best)
def main():
global args, best_mIoU
args = parser.parse_args()
os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(
str(gpu) for gpu in args.gpus)
args.gpus = len(args.gpus)
if args.no_partialbn:
sync_bn.Synchronize.init(args.gpus)
if args.dataset == 'VOCAug' or args.dataset == 'VOC2012' or args.dataset == 'COCO':
num_class = 21
ignore_label = 255
scale_series = [10, 20, 30, 60]
elif args.dataset == 'Cityscapes':
num_class = 19
ignore_label = 255 # 0
scale_series = [15, 30, 45, 90]
elif args.dataset == 'ApolloScape':
num_class = 37 # merge the noise and ignore labels
ignore_label = 255 # 0
else:
raise ValueError('Unknown dataset ' + args.dataset)
model = models.ERFNet(
num_class, partial_bn=not args.no_partialbn
) # models.PSPNet(num_class, base_model=args.arch, dropout=args.dropout, partial_bn=not args.no_partialbn)
input_mean = model.input_mean
input_std = model.input_std
policies = model.get_optim_policies()
model = torch.nn.DataParallel(model, device_ids=range(args.gpus)).cuda()
if args.resume:
if os.path.isfile(args.resume):
print(("=> loading checkpoint '{}'".format(args.resume)))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
best_mIoU = checkpoint['best_mIoU']
torch.nn.Module.load_state_dict(model, checkpoint['state_dict'])
print(("=> loaded checkpoint '{}' (epoch {})".format(
args.evaluate, checkpoint['epoch'])))
else:
print(("=> no checkpoint found at '{}'".format(args.resume)))
cudnn.benchmark = True
cudnn.fastest = True
# Data loading code
test_loader = torch.utils.data.DataLoader(getattr(
ds,
args.dataset.replace("ApolloScape", "VOCAug") + 'DataSet')(
data_list=args.val_list,
transform=[
torchvision.transforms.Compose([
tf.GroupRandomScaleRatio(
size=(1692, 1692, 505, 505),
interpolation=(cv2.INTER_LINEAR, cv2.INTER_NEAREST)),
tf.GroupNormalize(mean=(input_mean, (0, )),
std=(input_std, (1, ))),
]),
torchvision.transforms.Compose([
tf.GroupRandomScaleRatio(
size=(1861, 1861, 556, 556),
interpolation=(cv2.INTER_LINEAR, cv2.INTER_NEAREST)),
tf.GroupNormalize(mean=(input_mean, (0, )),
std=(input_std, (1, ))),
]),
torchvision.transforms.Compose([
tf.GroupRandomScaleRatio(
size=(1624, 1624, 485, 485),
interpolation=(cv2.INTER_LINEAR, cv2.INTER_NEAREST)),
tf.GroupNormalize(mean=(input_mean, (0, )),
std=(input_std, (1, ))),
]),
torchvision.transforms.Compose([
tf.GroupRandomScaleRatio(
size=(2030, 2030, 606, 606),
interpolation=(cv2.INTER_LINEAR, cv2.INTER_NEAREST)),
tf.GroupNormalize(mean=(input_mean, (0, )),
std=(input_std, (1, ))),
])
]),
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=False)
# define loss function (criterion) optimizer and evaluator
weights = [1.0 for _ in range(37)]
weights[0] = 0.05
weights[36] = 0.05
class_weights = torch.FloatTensor(weights).cuda()
criterion = torch.nn.NLLLoss(ignore_index=ignore_label,
weight=class_weights).cuda()
for group in policies:
print(('group: {} has {} params, lr_mult: {}, decay_mult: {}'.format(
group['name'], len(group['params']), group['lr_mult'],
group['decay_mult'])))
optimizer = torch.optim.SGD(policies,
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
evaluator = EvalSegmentation(num_class, ignore_label)
### evaluate ###
validate(test_loader, model, criterion, 0, evaluator)
return
def main():
global best_mIoU_cls, best_mIoU_ego, args
args = parser.parse_args()
os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(str(gpu) for gpu in cfg.gpus)
if cfg.dataset == 'VOCAug' or cfg.dataset == 'VOC2012' or cfg.dataset == 'COCO':
num_ego = 21
num_class = 2
ignore_label = 255
elif cfg.dataset == 'Cityscapes':
num_ego = 19
num_class = 2
ignore_label = 255 # 0
elif cfg.dataset == 'ApolloScape':
num_ego = 37 # merge the noise and ignore labels
num_class = 2
ignore_label = 255
elif cfg.dataset == 'CULane':
num_ego = cfg.NUM_EGO
num_class = 2
ignore_label = 255
else:
num_ego = cfg.NUM_EGO
num_class = cfg.NUM_CLASSES
ignore_label = 255
print(json.dumps(cfg, sort_keys=True, indent=2))
model = net.ERFNet(num_class, num_ego)
model = torch.nn.DataParallel(model,
device_ids=range(len(cfg.gpus))).cuda()
if num_class:
print(("=> train '{}' model".format('lane_cls')))
if num_ego:
print(("=> train '{}' model".format('lane_ego')))
if cfg.optimizer == 'sgd':
optimizer = torch.optim.SGD(model.parameters(),
cfg.lr,
momentum=cfg.momentum,
weight_decay=cfg.weight_decay)
else:
optimizer = torch.optim.Adam(model.parameters(),
cfg.lr,
weight_decay=cfg.weight_decay)
resume_epoch = 0
if cfg.resume:
if os.path.isfile(cfg.resume):
print(("=> loading checkpoint '{}'".format(cfg.resume)))
checkpoint = torch.load(cfg.resume)
if cfg.finetune:
print('finetune from ', cfg.resume)
state_all = checkpoint['state_dict']
state_clip = {} # only use backbone parameters
for k, v in state_all.items():
if 'module' in k:
state_clip[k] = v
print(k)
model.load_state_dict(state_clip, strict=False)
else:
print('==> Resume model from ' + cfg.resume)
model.load_state_dict(checkpoint['state_dict'])
if 'optimizer' in checkpoint.keys():
optimizer.load_state_dict(checkpoint['optimizer'])
if 'epoch' in checkpoint.keys():
resume_epoch = int(checkpoint['epoch']) + 1
else:
print(("=> no checkpoint found at '{}'".format(cfg.resume)))
model.apply(weights_init)
else:
model.apply(weights_init)
# if cfg.resume:
# if os.path.isfile(cfg.resume):
# print(("=> loading checkpoint '{}'".format(cfg.resume)))
# checkpoint = torch.load(cfg.resume)
# cfg.start_epoch = checkpoint['epoch']
# # model = load_my_state_dict(model, checkpoint['state_dict'])
# torch.nn.Module.load_state_dict(model, checkpoint['state_dict'])
# print(("=> loaded checkpoint '{}' (epoch {})".format(cfg.evaluate, checkpoint['epoch'])))
# else:
# print(("=> no checkpoint found at '{}'".format(cfg.resume)))
# model.apply(weights_init)
# else:
# model.apply(weights_init)
cudnn.benchmark = True
cudnn.fastest = True
# Data loading code
train_loader = torch.utils.data.DataLoader(getattr(ds, 'VOCAugDataSet')(
dataset_path=cfg.dataset_path,
data_list=cfg.train_list,
transform=torchvision.transforms.Compose([
tf.GroupRandomScale(size=(0.695, 0.721),
interpolation=(cv2.INTER_LINEAR,
cv2.INTER_NEAREST,
cv2.INTER_NEAREST)),
tf.GroupRandomCropRatio(size=(cfg.MODEL_INPUT_WIDTH,
cfg.MODEL_INPUT_HEIGHT)),
tf.GroupRandomRotation(degree=(-1, 1),
interpolation=(cv2.INTER_LINEAR,
cv2.INTER_NEAREST,
cv2.INTER_NEAREST),
padding=(cfg.INPUT_MEAN, (ignore_label, ),
(ignore_label, ))),
tf.GroupNormalize(mean=(cfg.INPUT_MEAN, (0, ), (0, )),
std=(cfg.INPUT_STD, (1, ), (1, ))),
])),
batch_size=cfg.train_batch_size,
shuffle=True,
num_workers=cfg.workers,
pin_memory=True,
drop_last=True)
val_loader = torch.utils.data.DataLoader(getattr(ds, 'VOCAugDataSet')(
dataset_path=cfg.dataset_path,
data_list=cfg.val_list,
transform=torchvision.transforms.Compose([
tf.GroupRandomScale(size=(0.695, 0.721),
interpolation=(cv2.INTER_LINEAR,
cv2.INTER_NEAREST,
cv2.INTER_NEAREST)),
tf.GroupRandomCropRatio(size=(cfg.MODEL_INPUT_WIDTH,
cfg.MODEL_INPUT_HEIGHT)),
tf.GroupNormalize(mean=(cfg.INPUT_MEAN, (0, ), (0, )),
std=(cfg.INPUT_STD, (1, ), (1, ))),
])),
batch_size=cfg.val_batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
# define loss function (criterion) optimizer and evaluator
class_weights = torch.FloatTensor(cfg.CLASS_WEIGHT).cuda()
weights = [1.0 for _ in range(num_ego + 1)]
weights[0] = 0.4
ego_weights = torch.FloatTensor(weights).cuda()
criterion_cls = torch.nn.NLLLoss(ignore_index=ignore_label,
weight=class_weights).cuda()
criterion_ego = torch.nn.NLLLoss(ignore_index=ignore_label,
weight=ego_weights).cuda()
criterion_exist = torch.nn.BCELoss().cuda()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
reg_loss = None
if cfg.weight_decay > 0 and cfg.use_L1:
reg_loss = Regularization(model, cfg.weight_decay, p=1).to(device)
else:
print("no regularization")
if num_class:
evaluator = EvalSegmentation(num_class, ignore_label)
if num_ego:
evaluator = EvalSegmentation(num_ego + 1, ignore_label)
# Tensorboard writer
global writer
writer = SummaryWriter(os.path.join(cfg.save_path, 'Tensorboard'))
for epoch in range(cfg.epochs): # args.start_epoch
if epoch < resume_epoch:
continue
adjust_learning_rate(optimizer, epoch, cfg.lr_steps)
# train for one epoch
train(train_loader, model, criterion_cls, criterion_ego,
criterion_exist, optimizer, epoch, writer, reg_loss)
# evaluate on validation set
if (epoch + 1) % cfg.eval_freq == 0 or epoch == cfg.epochs - 1:
mIoU_cls, mIoU_ego = validate(val_loader, model, criterion_cls,
criterion_ego, criterion_exist,
epoch, evaluator, writer)
# remember best mIoU and save checkpoint
if num_class:
is_best = mIoU_cls > best_mIoU_cls
if num_ego:
is_best = mIoU_ego > best_mIoU_ego
best_mIoU_cls = max(mIoU_cls, best_mIoU_cls)
best_mIoU_ego = max(mIoU_ego, best_mIoU_ego)
save_checkpoint(
{
'epoch': epoch + 1,
'arch': cfg.arch,
'state_dict': model.state_dict(),
'best_mIoU': best_mIoU_ego,
}, is_best)
writer.close()
def main():
global best_mIoU, start_epoch
if not os.path.exists(args.save_dir):
os.makedirs(args.save_dir)
if args.dataset == 'LaneDet':
num_class = 20
ignore_label = 255
else:
raise ValueError('Unknown dataset ' + args.dataset)
# get places
places = fluid.cuda_places()
with fluid.dygraph.guard():
model = models.ERFNet(num_class, [args.img_height, args.img_width])
input_mean = model.input_mean
input_std = model.input_std
# Data loading code
train_dataset = ds.LaneDataSet(
dataset_path='datasets/PreliminaryData',
data_list=args.train_list,
transform=[
tf.GroupRandomScale(size=(int(args.img_width), int(args.img_width * 1.2)),
interpolation=(cv2.INTER_LINEAR, cv2.INTER_NEAREST)),
tf.GroupRandomCropRatio(size=(args.img_width, args.img_height)),
tf.GroupNormalize(mean=(input_mean, (0,)), std=(input_std, (1,))),
]
)
train_loader = DataLoader(
train_dataset,
places=places[0],
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
drop_last=True
)
val_dataset = ds.LaneDataSet(
dataset_path='datasets/PreliminaryData',
data_list=args.train_list,
transform=[
tf.GroupRandomScale(size=args.img_width, interpolation=(cv2.INTER_LINEAR, cv2.INTER_NEAREST)),
tf.GroupNormalize(mean=(input_mean, (0,)), std=(input_std, (1,))),
],
is_val=False
)
val_loader = DataLoader(
val_dataset,
places=places[0],
batch_size=1,
shuffle=False,
num_workers=args.workers,
)
# define loss function (criterion) optimizer and evaluator
weights = [1.0 for _ in range(num_class)]
weights[0] = 0.25
weights = fluid.dygraph.to_variable(np.array(weights, dtype=np.float32))
criterion = fluid.dygraph.NLLLoss(weight=weights, ignore_index=ignore_label)
evaluator = EvalSegmentation(num_class, ignore_label)
optimizer = fluid.optimizer.MomentumOptimizer(learning_rate=fluid.dygraph.CosineDecay(
args.lr, len(train_loader), args.epochs),
momentum=args.momentum,
parameter_list=model.parameters(),
regularization=fluid.regularizer.L2Decay(
regularization_coeff=args.weight_decay))
if args.resume:
print(("=> loading checkpoint '{}'".format(args.resume)))
start_epoch = int(''.join([x for x in args.resume.split('/')[-1] if x.isdigit()]))
checkpoint, optim_checkpoint = fluid.load_dygraph(args.resume)
model.load_dict(checkpoint)
optimizer.set_dict(optim_checkpoint)
print(("=> loaded checkpoint (epoch {})".format(start_epoch)))
else:
try:
checkpoint, _ = fluid.load_dygraph(args.weight)
model.load_dict(checkpoint)
print("=> pretrained model loaded successfully")
except:
print(("=> no pretrained model found at '{}'".format(args.weight)))
for epoch in range(start_epoch, args.epochs):
# train for one epoch
loss = train(train_loader, model, criterion, optimizer, epoch)
# writer.add_scalar('lr', optimizer.current_step_lr(), epoch + 1)
if (epoch + 1) % args.save_freq == 0 or epoch == args.epochs - 1:
save_checkpoint(model.state_dict(), epoch)
save_checkpoint(optimizer.state_dict(), epoch)
# evaluate on validation set
if (epoch + 1) % args.eval_freq == 0 or epoch == args.epochs - 1:
mIoU = validate(val_loader, model, evaluator, epoch)
# remember best mIoU
is_best = mIoU > best_mIoU
best_mIoU = max(mIoU, best_mIoU)
if is_best:
tag_best(epoch, best_mIoU)