def get_dataset_interface(dataset_name, data_dir):
"""
Static function which returns an interface to a dataset.
"""
if dataset_name == "stanford":
return StanfordPedestrianData(data_dir)
elif dataset_name == "voc2012":
return VOCData(data_dir)
else:
raise "Unknown dataset: {}".format(dataset_name)
elif args.model == 'decoupled':
args.batch_size = 38
model_path = model_path + '.pth'
args.batch_size = 1
# net = net_saliency.resnet_saliency(flag_classify)
# net = net_saliency.Vgg16_for_saliency(flag_classify)
# net = net_saliency.SEC_for_saliency(flag_classify, args)
net = net_saliency.decouple_net_saliency(flag_classify)
print(args)
if flag_use_cuda:
net.cuda()
dataloader = VOCData(args)
net.train(False)
with torch.no_grad():
start = time.time()
for phase in ['train', 'val']:
if phase == 'train':
for data in dataloader.dataloaders["train"]:
inputs, labels, mask_gt, img = data
if flag_use_cuda:
inputs = inputs.cuda()
labels = labels.cuda()
outputs = net(inputs)