def main():
"""
Main Function
"""
# Set up the Arguments, Tensorboard Writer, Dataloader, Loss Fn, Optimizer
assert_and_infer_cfg(args)
writer = prep_experiment(args, parser)
train_loader, val_loader, train_obj = datasets.setup_loaders(args)
criterion, criterion_val = loss.get_loss(args)
net = network.get_net(args, criterion)
optim, scheduler = optimizer.get_optimizer(args, net)
if args.fix_bn:
net.apply(set_bn_eval)
print("Fix bn for finetuning")
if args.fp16:
net, optim = amp.initialize(net, optim, opt_level="O1")
net = network.wrap_network_in_dataparallel(net, args.apex)
if args.snapshot:
optimizer.load_weights(net, optim,
args.snapshot, args.restore_optimizer)
if args.evaluateF:
assert args.snapshot is not None, "must load weights for evaluation"
evaluate(val_loader, net, args)
return
torch.cuda.empty_cache()
# Main Loop
for epoch in range(args.start_epoch, args.max_epoch):
# Update EPOCH CTR
cfg.immutable(False)
cfg.EPOCH = epoch
cfg.immutable(True)
scheduler.step()
train(train_loader, net, optim, epoch, writer)
if args.apex:
train_loader.sampler.set_epoch(epoch + 1)
validate(val_loader, net, criterion_val,
optim, epoch, writer)
if args.class_uniform_pct:
if epoch >= args.max_cu_epoch:
train_obj.build_epoch(cut=True)
if args.apex:
train_loader.sampler.set_num_samples()
else:
train_obj.build_epoch()
def main():
"""
Main Function
"""
if AutoResume:
AutoResume.init()
assert args.result_dir is not None, 'need to define result_dir arg'
logx.initialize(logdir=args.result_dir,
tensorboard=True,
hparams=vars(args),
global_rank=args.global_rank)
# Set up the Arguments, Tensorboard Writer, Dataloader, Loss Fn, Optimizer
assert_and_infer_cfg(args)
prep_experiment(args)
train_loader, val_loader, train_obj = \
datasets.setup_loaders(args)
criterion, criterion_val = get_loss(args)
auto_resume_details = None
if AutoResume:
auto_resume_details = AutoResume.get_resume_details()
if auto_resume_details:
checkpoint_fn = auto_resume_details.get("RESUME_FILE", None)
checkpoint = torch.load(checkpoint_fn,
map_location=torch.device('cpu'))
args.result_dir = auto_resume_details.get("TENSORBOARD_DIR", None)
args.start_epoch = int(auto_resume_details.get("EPOCH", None)) + 1
args.restore_net = True
args.restore_optimizer = True
msg = ("Found details of a requested auto-resume: checkpoint={}"
" tensorboard={} at epoch {}")
logx.msg(msg.format(checkpoint_fn, args.result_dir, args.start_epoch))
elif args.resume:
checkpoint = torch.load(args.resume, map_location=torch.device('cpu'))
args.arch = checkpoint['arch']
args.start_epoch = int(checkpoint['epoch']) + 1
args.restore_net = True
args.restore_optimizer = True
msg = "Resuming from: checkpoint={}, epoch {}, arch {}"
logx.msg(msg.format(args.resume, args.start_epoch, args.arch))
elif args.snapshot:
if 'ASSETS_PATH' in args.snapshot:
args.snapshot = args.snapshot.replace('ASSETS_PATH',
cfg.ASSETS_PATH)
checkpoint = torch.load(args.snapshot,
map_location=torch.device('cpu'))
args.restore_net = True
msg = "Loading weights from: checkpoint={}".format(args.snapshot)
logx.msg(msg)
#define the NASA optimizer parameter
iter_tot = len(train_loader) * args.max_epoch
# tau = args.tau_factor/sqrt(iter_tot)
tau = 1
net = network.get_net(args, criterion)
k = 1
# optim, scheduler = get_optimizer(args, net)
optim, scheduler = get_optimizer(args, net, tau, k)
# Visualize feature maps
#activation = {}
#def get_activation(name):
#def hook(model, input, output):
#activation[name] = output.detach()
#return hook
#net.layer[0].register_forward_hook(get_activation('conv1'))
#data, _ = dataset[0]
#data.unsqueeze_(0)
#output = model(data)
#act = activation['conv1'].squeeze()
#fig, axarr = plt.subplots(act.size(0))
#for idx in range(act.size(0)):
#axarr[idx].imshow(act[idx])
if args.fp16:
net, optim = amp.initialize(net, optim, opt_level=args.amp_opt_level)
net = network.wrap_network_in_dataparallel(net, args.apex)
if args.summary:
from thop import profile
img = torch.randn(1, 3, 640, 640).cuda()
mask = torch.randn(1, 1, 640, 640).cuda()
macs, params = profile(net, inputs={'images': img, 'gts': mask})
print(f'macs {macs} params {params}')
sys.exit()
if args.restore_optimizer:
restore_opt(optim, checkpoint)
if args.restore_net:
restore_net(net, checkpoint)
if args.init_decoder:
net.module.init_mods()
torch.cuda.empty_cache()
if args.start_epoch != 0:
scheduler.step(args.start_epoch)
# There are 4 options for evaluation:
# --eval val just run validation
# --eval val --dump_assets dump all images and assets
# --eval folder just dump all basic images
# --eval folder --dump_assets dump all images and assets
if args.eval == 'test':
validate(val_loader,
net,
criterion=None,
optim=None,
epoch=0,
calc_metrics=False,
dump_assets=args.dump_assets,
dump_all_images=True,
testing=True,
grid=city)
return 0
if args.eval == 'val':
if args.dump_topn:
validate_topn(val_loader, net, criterion_val, optim, 0, args)
else:
validate(val_loader,
net,
criterion=criterion_val,
optim=optim,
epoch=0,
dump_assets=args.dump_assets,
dump_all_images=args.dump_all_images,
calc_metrics=not args.no_metrics)
return 0
elif args.eval == 'folder':
# Using a folder for evaluation means to not calculate metrics
validate(val_loader,
net,
criterion=criterion_val,
optim=optim,
epoch=0,
calc_metrics=False,
dump_assets=args.dump_assets,
dump_all_images=True)
return 0
elif args.eval is not None:
raise 'unknown eval option {}'.format(args.eval)
for epoch in range(args.start_epoch, args.max_epoch):
update_epoch(epoch)
if args.only_coarse:
train_obj.only_coarse()
train_obj.build_epoch()
if args.apex:
train_loader.sampler.set_num_samples()
elif args.class_uniform_pct:
if epoch >= args.max_cu_epoch:
train_obj.disable_coarse()
train_obj.build_epoch()
if args.apex:
train_loader.sampler.set_num_samples()
else:
train_obj.build_epoch()
else:
pass
train(train_loader, net, optim, epoch)
if args.apex:
train_loader.sampler.set_epoch(epoch + 1)
if epoch % args.val_freq == 0:
validate(val_loader, net, criterion_val, optim, epoch)
scheduler.step()
if check_termination(epoch):
return 0
def main():
"""
Main Function
"""
if AutoResume:
AutoResume.init()
assert args.result_dir is not None, 'need to define result_dir arg'
logx.initialize(logdir=args.result_dir,
tensorboard=True, hparams=vars(args),
global_rank=args.global_rank)
# Set up the Arguments, Tensorboard Writer, Dataloader, Loss Fn, Optimizer
assert_and_infer_cfg(args)
prep_experiment(args)
train_loader, val_loader, train_obj = \
datasets.setup_loaders(args)
criterion, criterion_val = get_loss(args)
auto_resume_details = None
if AutoResume:
auto_resume_details = AutoResume.get_resume_details()
if auto_resume_details:
checkpoint_fn = auto_resume_details.get("RESUME_FILE", None)
checkpoint = torch.load(checkpoint_fn,
map_location=torch.device('cpu'))
args.result_dir = auto_resume_details.get("TENSORBOARD_DIR", None)
args.start_epoch = int(auto_resume_details.get("EPOCH", None)) + 1
args.restore_net = True
args.restore_optimizer = True
msg = ("Found details of a requested auto-resume: checkpoint={}"
" tensorboard={} at epoch {}")
logx.msg(msg.format(checkpoint_fn, args.result_dir,
args.start_epoch))
elif args.resume:
checkpoint = torch.load(args.resume,
map_location=torch.device('cpu'))
args.arch = checkpoint['arch']
args.start_epoch = int(checkpoint['epoch']) + 1
args.restore_net = True
args.restore_optimizer = True
msg = "Resuming from: checkpoint={}, epoch {}, arch {}"
logx.msg(msg.format(args.resume, args.start_epoch, args.arch))
elif args.snapshot:
if 'ASSETS_PATH' in args.snapshot:
args.snapshot = args.snapshot.replace('ASSETS_PATH', cfg.ASSETS_PATH)
checkpoint = torch.load(args.snapshot,
map_location=torch.device('cpu'))
args.restore_net = True
msg = "Loading weights from: checkpoint={}".format(args.snapshot)
logx.msg(msg)
net = network.get_net(args, criterion)
optim, scheduler = get_optimizer(args, net)
if args.fp16:
net, optim = amp.initialize(net, optim, opt_level=args.amp_opt_level)
net = network.wrap_network_in_dataparallel(net, args.apex)
if args.summary:
print(str(net))
from pytorchOpCounter.thop import profile
img = torch.randn(1, 3, 1024, 2048).cuda()
mask = torch.randn(1, 1, 1024, 2048).cuda()
macs, params = profile(net, inputs={'images': img, 'gts': mask})
print(f'macs {macs} params {params}')
sys.exit()
if args.restore_optimizer:
restore_opt(optim, checkpoint)
if args.restore_net:
restore_net(net, checkpoint)
if args.init_decoder:
net.module.init_mods()
torch.cuda.empty_cache()
if args.start_epoch != 0:
scheduler.step(args.start_epoch)
# There are 4 options for evaluation:
# --eval val just run validation
# --eval val --dump_assets dump all images and assets
# --eval folder just dump all basic images
# --eval folder --dump_assets dump all images and assets
if args.eval == 'val':
if args.dump_topn:
validate_topn(val_loader, net, criterion_val, optim, 0, args)
else:
validate(val_loader, net, criterion=criterion_val, optim=optim, epoch=0,
dump_assets=args.dump_assets,
dump_all_images=args.dump_all_images,
calc_metrics=not args.no_metrics)
return 0
elif args.eval == 'folder':
# Using a folder for evaluation means to not calculate metrics
validate(val_loader, net, criterion=None, optim=None, epoch=0,
calc_metrics=False, dump_assets=args.dump_assets,
dump_all_images=True)
return 0
elif args.eval is not None:
raise 'unknown eval option {}'.format(args.eval)
for epoch in range(args.start_epoch, args.max_epoch):
update_epoch(epoch)
if args.only_coarse:
train_obj.only_coarse()
train_obj.build_epoch()
if args.apex:
train_loader.sampler.set_num_samples()
elif args.class_uniform_pct:
if epoch >= args.max_cu_epoch:
train_obj.disable_coarse()
train_obj.build_epoch()
if args.apex:
train_loader.sampler.set_num_samples()
else:
train_obj.build_epoch()
else:
pass
train(train_loader, net, optim, epoch)
if args.apex:
train_loader.sampler.set_epoch(epoch + 1)
if epoch % args.val_freq == 0:
validate(val_loader, net, criterion_val, optim, epoch)
scheduler.step()
if check_termination(epoch):
return 0
def main():
"""
Main Function
"""
if AutoResume:
AutoResume.init()
assert args.result_dir is not None, 'need to define result_dir arg'
logx.initialize(logdir=args.result_dir,
tensorboard=False,
hparams=vars(args),
global_rank=args.global_rank)
# Set up the Arguments, Tensorboard Writer, Dataloader, Loss Fn, Optimizer
assert_and_infer_cfg(args)
prep_experiment(args)
train_loader, val_loader, train_obj = datasets.setup_loaders(args)
criterion, criterion_val = get_loss(args)
auto_resume_details = None
if AutoResume:
auto_resume_details = AutoResume.get_resume_details()
if auto_resume_details:
checkpoint_fn = auto_resume_details.get("RESUME_FILE", None)
checkpoint = torch.load(checkpoint_fn,
map_location=torch.device('cpu'))
args.result_dir = auto_resume_details.get("TENSORBOARD_DIR", None)
args.start_epoch = int(auto_resume_details.get("EPOCH", None)) + 1
args.restore_net = True
args.restore_optimizer = True
msg = ("Found details of a requested auto-resume: checkpoint={}"
" tensorboard={} at epoch {}")
logx.msg(msg.format(checkpoint_fn, args.result_dir, args.start_epoch))
elif args.resume:
checkpoint = torch.load(args.resume, map_location=torch.device('cpu'))
args.arch = checkpoint['arch']
args.start_epoch = int(checkpoint['epoch']) + 1
args.restore_net = True
args.restore_optimizer = True
msg = "Resuming from: checkpoint={}, epoch {}, arch {}"
logx.msg(msg.format(args.resume, args.start_epoch, args.arch))
elif args.snapshot:
if 'ASSETS_PATH' in args.snapshot:
args.snapshot = args.snapshot.replace('ASSETS_PATH',
cfg.ASSETS_PATH)
checkpoint = torch.load(args.snapshot,
map_location=torch.device('cpu'))
args.restore_net = True
msg = "Loading weights from: checkpoint={}".format(args.snapshot)
logx.msg(msg)
net = network.get_net(args, criterion)
optim, scheduler = get_optimizer(args, net)
net = network.wrap_network_in_dataparallel(net, args.apex)
if args.restore_optimizer:
restore_opt(optim, checkpoint)
if args.restore_net:
restore_net(net, checkpoint)
if args.init_decoder:
net.module.init_mods()
torch.cuda.empty_cache()
if args.start_epoch != 0:
scheduler.step(args.start_epoch)
if args.eval == 'folder':
# Using a folder for evaluation means to not calculate metrics
# validate(val_loader, net, criterion=None, optim=None, epoch=0,
# calc_metrics=False, dump_assets=args.dump_assets,
# dump_all_images=True)
if not os.path.exists(args.result_dir + 'image_2/'):
os.mkdir(args.result_dir + 'image_2/')
if not os.path.exists(args.result_dir + 'image_3/'):
os.mkdir(args.result_dir + 'image_3/')
num_image = 7481
for idx in tqdm(range(num_image)):
sample_idx = "%06d" % idx
eval_minibatch(sample_idx, "image_2/", net, args)
eval_minibatch(sample_idx, "image_3/", net, args)
return 0
elif args.eval is not None:
raise 'unknown eval option {}'.format(args.eval)
def main():
"""
Main Function
"""
# Set up the Arguments, Tensorboard Writer, Dataloader, Loss Fn, Optimizer
# args2 = copy.deepcopy(args)
assert_and_infer_cfg(args)
# assert_and_infer_cfg(args2)
# args2.dataset = 'kitti_trav'
# print(args.dataset)
# print(args2.dataset)
writer = prep_experiment(args, parser)
# writer = prep_experiment(args2, parser)
# Dataset
train_loader, val_loader, train_obj = datasets.setup_loaders(args)
# train_loader2, val_loader2, train_obj2 = datasets.setup_loaders(args2)
criterion, criterion_val = loss.get_loss(args, data_type='semantic')
criterion2, criterion_val2 = loss.get_loss(args, data_type='trav')
net = network.get_net(args, criterion, criterion2)
#parameters list
# param1_lists = list(net.mod1.parameters()) + list(net.mod2.parameters()) + list(net.mod3.parameters()) + list(net.mod4.parameters()) + list(net.mod5.parameters()) + list(net.mod6.parameters()) + list(net.mod7.parameters()) + list(net.pool2.parameters()) + list(net.pool3.parameters()) + list(net.aspp.parameters()) + list(net.bot_fine.parameters()) + list(net.bot_aspp.parameters()) + list(net.final.parameters()) + [log_sigma_A]
# param2_lists = list(net.mod1.parameters()) + list(net.mod2.parameters()) + list(net.mod3.parameters()) + list(net.mod4.parameters()) + list(net.mod5.parameters()) + list(net.mod6.parameters()) + list(net.mod7.parameters()) + list(net.pool2.parameters()) + list(net.pool3.parameters()) + list(net.aspp.parameters()) + list(net.bot_fine.parameters()) + list(net.bot_aspp.parameters()) + list(net.final2.parameters()) + [log_sigma_B]
#optimizers
optim, scheduler = optimizer.get_optimizer(args, net)
# optim2, scheduler2 = optimizer.get_optimizer(args, param2_lists)
if args.fp16:
net, optim = amp.initialize(net, optim, opt_level="O1")
net = network.wrap_network_in_dataparallel(net, args.apex)
if args.snapshot:
optimizer.load_weights(net, optim, args.snapshot, args.snapshot2,
args.restore_optimizer)
# optimizer.load_weights(net, optim2,
# args.snapshot, args.snapshot2, args.restore_optimizer)
torch.cuda.empty_cache()
# Main Loop
for epoch in range(args.start_epoch, args.max_epoch):
# Update EPOCH CTR
cfg.immutable(False)
cfg.EPOCH = epoch
cfg.immutable(True)
scheduler.step()
train(train_loader, net, optim, epoch, writer)
if args.apex:
train_loader.sampler.set_epoch(epoch + 1)
# train_loader2.sampler.set_epoch(epoch + 1)
validate(val_loader, net, criterion_val, criterion_val2, optim, epoch,
writer)
if args.class_uniform_pct:
if epoch >= args.max_cu_epoch:
train_obj.build_epoch(cut=True)
# train_obj2.build_epoch(cut=True)
if args.apex:
train_loader.sampler.set_num_samples()
# train_loader2.sampler.set_num_samples()
else:
train_obj.build_epoch()
