def get_net():
"""
Get Network for evaluation
"""
logging.info('Load model file: %s', args.snapshot)
net = network.get_net_ori(args, criterion=None)
net = torch.nn.DataParallel(net).cuda()
net, _ = restore_snapshot(net, optimizer=None,
snapshot=args.snapshot, snapshot2=args.snapshot2, restore_optimizer_bool=False)
net.eval()
return net
def main():
"""
Main Function
"""
# Set up the Arguments, Tensorboard Writer, Dataloader, Loss Fn, Optimizer
assert_and_infer_cfg(args)
writer = prep_experiment_ori(args, parser)
train_loader, val_loader, train_obj = datasets.setup_loaders(args)
criterion, criterion_val = loss.get_loss(args, data_type = 'trav_alone')
net = network.get_net_ori(args, criterion)
optim, scheduler = optimizer.get_optimizer(args, net)
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)
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()
# set the device
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print('running on device ' + str(device))
# load the model checkpoint
print('loading checkpoint: ' + opt.input)
checkpoint = torch.load(opt.input)
num_classes = opt.num_classes
# create the model architecture
print('num classes: ' + str(num_classes))
criterion = CrossEntropyLoss2d(size_average=True, ignore_index=None).cuda()
model = network.get_net_ori(opt, criterion)
model = nn.DataParallel(model)
# load the model weights
model.load_state_dict(checkpoint)
model.to(device)
model.eval()
print(model)
print('')
# create example image data
resolution = [opt.width, opt.height]
input = torch.ones((1, 3, resolution[0], resolution[1])).cuda()
print('input size: {:d}x{:d}'.format(resolution[1], resolution[0]))