def convert_to_onnx(pth, length):
cd = path.dirname(path.abspath(__file__))
config_file = path.join(cd, 'config.json')
config = ConfigParser.from_json(config_file)
logger = config.get_logger('convert', 1)
model = config.init_obj('arch', module_arch)
state = torch.load(pth, map_location=torch.device('cpu'))
state_dict = state['state_dict']
model.load_state_dict(state_dict)
model.eval()
vocabulary = list('ACGT')
encodes = np.eye(len(vocabulary))
x = encodes[np.random.choice(encodes.shape[0], size=length)]
input_tensor = torch.FloatTensor(np.expand_dims(x, axis=0))
exported_onnx_file = pth.rsplit('.', 1)[0] + '.onnx'
logger.info('Converting to ONNX model: {}'.format(exported_onnx_file))
with warnings.catch_warnings():
warnings.simplefilter("ignore")
torch.onnx.export(model,
input_tensor,
exported_onnx_file,
export_params=True,
opset_version=10,
do_constant_folding=True,
input_names=["input"],
output_names=["output"],
dynamic_axes={
"input": {
0: "batch_size",
1: "sequence"
},
"output": {
0: "batch_size",
1: "sequence"
}
})
args.add_argument(
'--chunk_size',
default=1024,
type=int,
help=
'chunk_size * threads reads to process per thread.(default: 1024) \n{}.'
.format(
'When chunk_size=1024 and threads=20, each process will load 1024 reads, in total consumming ~20G memory'
))
args.add_argument('-v',
'--version',
action='version',
version='%(prog)s {version}'.format(version=__version__))
if not isinstance(args, tuple):
args = args.parse_args()
if args.config is None:
config_file = os.path.join(cd, 'config.json')
else:
config_file = args.config
config = ConfigParser.from_json(config_file)
os.environ['OMP_NUM_THREADS'] = '1'
# os.environ['MKL_NUM_THREADS'] = '1'
seq_pred = Predictor(config, args)
seq_pred.load_model()
seq_pred.run()
