def build_caffe2_model_file(model, file_name, inputs_spec: dict, output_spec: dict): with open(CAFFE2_MODEL_PREDICT_FILE_NAME, 'wb') as f: f.write(model.net._net.SerializeToString()) init_net = caffe2_pb2.NetDef() for param in model.params: op = core.CreateOperator( "GivenTensorFill", [], [param], arg=[ utils.MakeArgument("shape", workspace.FetchBlob(param).shape), utils.MakeArgument("values", workspace.FetchBlob(param)) ] ) init_net.op.extend([op]) init_net.op.extend([core.CreateOperator( "ConstantFill", [], [inputs_spec['inputs']['0']['name']], shape=tuple(inputs_spec['inputs']['0']['shape']) )] ) with open(CAFFE2_MODEL_INIT_FILE_NAME, 'wb') as f: f.write(init_net.SerializeToString()) FileHelper.write_to_file(INPUT_SPEC_FILE_NAME, json.dumps(inputs_spec)) FileHelper.write_to_file(OUTPUT_SPEC_FILE_NAME, json.dumps(output_spec)) FileHelper.write_files_to_tar( file_name + '.caffe2', SerializationHelper.get_list_of_model_file_content('caffe2') ) os.remove(CAFFE2_MODEL_INIT_FILE_NAME) os.remove(CAFFE2_MODEL_PREDICT_FILE_NAME) os.remove(INPUT_SPEC_FILE_NAME) os.remove(OUTPUT_SPEC_FILE_NAME)
def build_sklearn_model_file(model, file_name, inputs_spec: dict, output_spec: dict): joblib.dump(model, SKLEARN_MODEL_FILE_NAME) FileHelper.write_to_file(INPUT_SPEC_FILE_NAME, json.dumps(inputs_spec)) FileHelper.write_to_file(OUTPUT_SPEC_FILE_NAME, json.dumps(output_spec)) FileHelper.write_files_to_tar( file_name + '.sklearn', SerializationHelper.get_list_of_model_file_content('sklearn') ) os.remove(SKLEARN_MODEL_FILE_NAME) os.remove(INPUT_SPEC_FILE_NAME) os.remove(OUTPUT_SPEC_FILE_NAME)