def main(args):
# read all the images in the folder
image_list = glob.glob(args.data_dir + os.sep + '*.' + args.img_extn)
up = None
if args.modelType == 2:
up = torch.nn.Upsample(scale_factor=16, mode='bilinear')
else:
up = torch.nn.Upsample(scale_factor=2, mode='bilinear')
if args.gpu:
up = up.cuda()
p = args.p
q = args.q
classes = args.classes
if args.modelType == 2:
modelA = Net.ESPNet_Encoder(
classes, p, q) # Net.Mobile_SegNetDilatedIA_C_stage1(20)
model_weight_file = args.weightsDir + os.sep + 'encoder' + os.sep + 'espnet_p_' + str(
p) + '_q_' + str(q) + '.pth'
if not os.path.isfile(model_weight_file):
print(
'Pre-trained model file does not exist. Please check ../pretrained/encoder folder'
)
exit(-1)
modelA.load_state_dict(torch.load(model_weight_file))
elif args.modelType == 1:
modelA = Net.ESPNet(classes, p,
q) # Net.Mobile_SegNetDilatedIA_C_stage1(20)
model_weight_file = args.weightsDir + os.sep + 'decoder' + os.sep + 'espnet_p_' + str(
p) + '_q_' + str(q) + '.pth'
if not os.path.isfile(model_weight_file):
print(
'Pre-trained model file does not exist. Please check ../pretrained/decoder folder'
)
exit(-1)
modelA.load_state_dict(torch.load(model_weight_file))
else:
print('Model not supported')
# modelA = torch.nn.DataParallel(modelA)
if args.gpu:
modelA = modelA.cuda()
# set to evaluation mode
modelA.eval()
if not os.path.isdir(args.savedir):
os.mkdir(args.savedir)
evaluateModel(args, modelA, up, image_list)
def main(args):
classes = args.classes
p = args.p
q = args.q
if args.modelType == 2:
model = Net.ESPNet_Encoder(
classes, p, q) # Net.Mobile_SegNetDilatedIA_C_stage1(20)
model_weight_file = args.weightsDir + os.sep + 'encoder' + os.sep + 'espnet_p_' + str(
p) + '_q_' + str(q) + '.pth'
if not os.path.isfile(model_weight_file):
print(
'Pre-trained model file does not exist. Please check ../pretrained/encoder folder'
)
exit(-1)
model.load_state_dict(torch.load(model_weight_file))
elif args.modelType == 1:
model = Net.ESPNet(classes, p,
q) # Net.Mobile_SegNetDilatedIA_C_stage1(20)
model_weight_file = args.weightsDir + os.sep + 'decoder' + os.sep + 'espnet_p_' + str(
p) + '_q_' + str(q) + '.pth'
if not os.path.isfile(model_weight_file):
print(
'Pre-trained model file does not exist. Please check ../pretrained/decoder folder'
)
exit(-1)
model.load_state_dict(torch.load(model_weight_file))
else:
print('Model not supported')
if (not args.cpu):
model = model.cuda(
) #.half() #HALF seems to be doing slower for some reason
#model = torch.nn.DataParallel(model).cuda()
model.eval()
images = torch.randn(args.batch_size, args.num_channels, args.height,
args.width)
if (not args.cpu):
images = images.cuda() #.half()
time_train = []
i = 0
while (True):
#for step, (images, labels, filename, filenameGt) in enumerate(loader):
start_time = time.time()
inputs = Variable(images, volatile=True)
outputs = model(inputs)
#preds = outputs.cpu()
if (not args.cpu):
torch.cuda.synchronize(
) #wait for cuda to finish (cuda is asynchronous!)
if i != 0: #first run always takes some time for setup
fwt = time.time() - start_time
time_train.append(fwt)
print("Forward time per img (b=%d): %.3f (Mean: %.3f)" %
(args.batch_size, fwt / args.batch_size,
sum(time_train) / len(time_train) / args.batch_size))
time.sleep(1) #to avoid overheating the GPU too much
i += 1
