def main():
logger.info(args)
model = Net(numAngle=CONFIGS["MODEL"]["NUMANGLE"],
numRho=CONFIGS["MODEL"]["NUMRHO"],
backbone=CONFIGS["MODEL"]["BACKBONE"])
if args.model:
if isfile(args.model):
import torch
m = torch.load(args.model)
if 'state_dict' in m.keys():
m = m['state_dict']
torch.save(m, '_temp_model.pth')
del m
logger.info("=> loading pretrained model '{}'".format(args.model))
#model.load('_temp_model.pth')
logger.info("=> loaded checkpoint '{}'".format(args.model))
else:
logger.info("=> no pretrained model found at '{}'".format(
args.model))
# dataloader
test_loader = get_loader(CONFIGS["DATA"]["TEST_DIR"],
CONFIGS["DATA"]["TEST_LABEL_FILE"],
batch_size=int(os.environ.get("BS", "1")),
num_thread=CONFIGS["DATA"]["WORKERS"],
test=True)
logger.info("Data loading done.")
weights_nodes = {}
data_nodes = {}
def named_dump_func(name):
def dump_func(self, inputs, outputs):
input_name = name + '_input'
output_name = name + '_output'
if isinstance(self, nn.Conv2d):
weights_nodes[name] = self.weight.numpy()
data_nodes[input_name] = inputs[0].numpy()
data_nodes[output_name] = outputs[0].numpy()
return dump_func
if args.dump:
logger.info('Add hooks to dump data.')
for name, module in model.named_modules():
print(name)
module.register_forward_hook(named_dump_func(name))
test(test_loader, model, args)
def main():
logger.info(args)
model = Net(numAngle=CONFIGS["MODEL"]["NUMANGLE"], numRho=CONFIGS["MODEL"]["NUMRHO"], backbone=CONFIGS["MODEL"]["BACKBONE"])
if args.model:
if isfile(args.model):
logger.info("=> loading pretrained model '{}'".format(args.model))
import torch
m = torch.load(args.model)
if 'state_dict' in m.keys():
m = m['state_dict']
torch.save(m, '_temp_model.pth')
del m
model.load('_temp_model.pth')
logger.info("=> loaded checkpoint '{}'".format(args.model))
else:
logger.info("=> no pretrained model found at '{}'".format(args.model))
# dataloader
test_loader = get_loader(CONFIGS["DATA"]["TEST_DIR"], CONFIGS["DATA"]["TEST_LABEL_FILE"],
batch_size=1, num_thread=CONFIGS["DATA"]["WORKERS"], test=True)
logger.info("Data loading done.")
weights_nodes = {}
data_nodes = {}
def named_dump_func(name):
def dump_func(self, inputs, outputs):
input_name = name + '_input'
output_name = name + '_output'
if isinstance(self, nn.Conv2d):
weights_nodes[name] = self.weight.numpy()
data_nodes[input_name] = inputs[0].numpy()
data_nodes[output_name] = outputs[0].numpy()
return dump_func
if args.dump:
logger.info('Add hooks to dump data.')
for name, module in model.named_modules():
module.register_forward_hook(named_dump_func(name))
logger.info("Start testing.")
total_time = test(test_loader, model, args)
if args.dump:
np.save('data_nodes.npy', data_nodes)
np.save('weights_nodes.npy', weights_nodes)
exit()
logger.info("Test done! Total %d imgs at %.4f secs without image io, fps: %.3f" % (len(test_loader), total_time, len(test_loader) / total_time))