def __init__(self, flag_classify, args):
super(SEC_for_saliency, self).__init__()
self.net = sec.SEC_NN(args.batch_size, args.num_classes,
args.output_size, args.no_bg, False)
self.net.load_state_dict(torch.load("models/sec_rename_CPU.pth"),
strict=False)
self.net.train(False)
self.classify = flag_classify
def __init__(self, flag_classify, args):
super(sec_for_saliency, self).__init__()
net = sec.SEC_NN(args.batch_size, args.num_classes, args.output_size,
args.no_bg, False)
#net.load_state_dict(torch.load("models/sec_rename_CPU.pth"), strict = True)
net.load_state_dict(torch.load("models/01/top_val_acc_SEC_CPU.pth"),
strict=True)
net.train(False)
self.classify = flag_classify
self.classifier = net.mask2label_pool
self.features = net.features
if self.classify:
self.classifier = net.mask2label_pool
elif args.model == 'my_resnet':
args.batch_size = 30
elif args.model == 'decoupled':
args.batch_size = 38
model_path = model_path + '.pth'
# if args.origin_size:
# args.batch_size = 1
if args.model == 'SEC':
# model_url = 'https://download.pytorch.org/models/vgg16-397923af.pth' # 'vgg16'
# model_path = 'models/0506/top_val_rec_SEC_05_CPU.pth' # 'vgg16'
args.input_size = [321, 321]
args.output_size = [41, 41]
net = sec.SEC_NN(args.batch_size, args.num_classes, args.output_size,
args.no_bg, flag_use_cuda)
net.load_state_dict(torch.load(model_path), strict=True)
elif args.model == 'resnet':
net = resnet.resnet50(pretrained=False, num_classes=args.num_classes)
net.load_state_dict(torch.load(model_path), strict=True)
features_blob = []
params = list(net.parameters())
fc_weight = params[-2]
def hook_feature(module, input, output):
features_blob.append(output.data)
net._modules.get('layer4').register_forward_hook(hook_feature)
elif args.model == 'my_resnet':