def get_tf_mod(verbose=False): """ :param bool verbose: :return: module """ global _tf_mod if _tf_mod: return _tf_mod import platform from glob import glob from returnn.tf.util.basic import OpCodeCompiler # References: # https://github.com/kpu/kenlm/blob/master/setup.py # https://github.com/kpu/kenlm/blob/master/compile_query_only.sh # Collect files. assert kenlm_checked_out(), "submodule in %r not checked out?" % kenlm_dir files = glob('%s/util/*.cc' % kenlm_dir) files += glob('%s/lm/*.cc' % kenlm_dir) files += glob('%s/util/double-conversion/*.cc' % kenlm_dir) files = [fn for fn in files if not (fn.endswith('main.cc') or fn.endswith('test.cc'))] assert files, "submodule in %r not checked out?" % kenlm_dir libs = ["z"] if platform.system() != 'Darwin': libs.append('rt') # Put code all together in one big blob. src_code = "" src_code += _kenlm_src_code_workarounds for fn in files: f_code = open(fn).read() f_code = ''.join([x for x in f_code if ord(x) < 128]) # enforce ASCII # We need to do some replacements to not clash symbol names. fn_short = os.path.basename(fn).replace(".", "_") for word in ["kConverter"]: f_code = f_code.replace(word, "%s_%s" % (fn_short, word)) src_code += "\n// ------------ %s : BEGIN { ------------\n" % os.path.basename(fn) # https://gcc.gnu.org/onlinedocs/cpp/Line-Control.html#Line-Control src_code += "#line 1 \"%s\"\n" % os.path.basename(fn) src_code += f_code src_code += "\n// ------------ %s : END } --------------\n\n" % os.path.basename(fn) src_code += "\n\n// ------------ our code now: ------------\n\n" src_code += "#line 1 \"<our code>\"\n" src_code += _src_code compiler = OpCodeCompiler( base_name="KenLM", code_version=1, code=src_code, include_paths=(kenlm_dir, kenlm_dir + "/util/double-conversion"), c_macro_defines={"NDEBUG": 1, "KENLM_MAX_ORDER": 6, "HAVE_ZLIB": 1}, ld_flags=["-l%s" % lib for lib in libs], is_cpp=True, use_cuda_if_available=False, verbose=verbose) tf_mod = compiler.load_tf_module() assert hasattr(tf_mod, "ken_lm_abs_score_strings"), "content of mod: %r" % (dir(tf_mod),) _tf_mod = tf_mod return tf_mod
def get_tf_mod(verbose=False): """ :param bool verbose: :return: module """ global _tf_mod if _tf_mod: return _tf_mod from glob import glob from returnn.tf.util.basic import OpCodeCompiler # Collect files. assert openfst_checked_out( ), "submodule in %r not checked out?" % openfst_dir files = glob('%s/src/lib/*.cc' % openfst_dir) assert files, "submodule in %r not checked out?" % openfst_dir files = sorted(files) # make somewhat deterministic libs = [] if platform.system() != 'Darwin': libs.append('rt') # Put code all together in one big blob. src_code = "" for fn in files: f_code = open(fn).read() f_code = ''.join([x for x in f_code if ord(x) < 128]) # enforce ASCII src_code += "\n// ------------ %s : BEGIN { ------------\n" % os.path.basename( fn) # https://gcc.gnu.org/onlinedocs/cpp/Line-Control.html#Line-Control src_code += "#line 1 \"%s\"\n" % os.path.basename(fn) src_code += f_code src_code += "\n// ------------ %s : END } --------------\n\n" % os.path.basename( fn) src_code += "\n\n// ------------ our code now: ------------\n\n" src_code += "#line 1 \"returnn/tf/util/open_fst.py:_src_code\"\n" src_code += _src_code compiler = OpCodeCompiler( base_name="OpenFst", code_version=1, code=src_code, include_paths=("%s/src/include" % openfst_dir, ), c_macro_defines={ "NDEBUG": 1, # https://github.com/tensorflow/tensorflow/issues/17316 }, ld_flags=["-l%s" % lib for lib in libs], is_cpp=True, use_cuda_if_available=False, verbose=verbose) tf_mod = compiler.load_tf_module() assert hasattr( tf_mod, "open_fst_transition"), "content of mod: %r" % (dir(tf_mod), ) _tf_mod = tf_mod return tf_mod
def init_warprnnt(verbose=False): """ Initializes and compiles the library. Caches the TF module. :param bool verbose: """ global _tf_mod if _tf_mod: return assert is_checked_out( ), "submodule not checked out? Run `git submodule update --init --recursive`" # References: # https://github.com/HawkAaron/warp-transducer/blob/master/tensorflow_binding/setup.py src_files = [ '%s/src/rnnt_entrypoint.cpp' % submodule_dir, '%s/tensorflow_binding/src/warprnnt_op.cc' % submodule_dir ] assert all([os.path.isfile(f) for f in src_files ]), "submodule in %r not checked out?" % warprnnt_dir src_code = "" for fn in src_files: f_code = open(fn).read() src_code += "\n// ------------ %s : BEGIN { ------------\n" % os.path.basename( fn) # https://gcc.gnu.org/onlinedocs/cpp/Line-Control.html#Line-Control src_code += "#line 1 \"%s\"\n" % os.path.basename(fn) src_code += f_code src_code += "\n// ------------ %s : END } --------------\n\n" % os.path.basename( fn) with_cuda = CudaEnv.get_instance().is_available() and is_gpu_available() with_omp = sys.platform.startswith("linux") compiler = OpCodeCompiler( base_name="warprnnt_kernels", code_version=1, code=src_code, include_paths=(submodule_dir + "/include", ), c_macro_defines=dict_joined( {"WITH_OMP": 1} if with_omp else {"RNNT_DISABLE_OMP": 1}, {"WARPRNNT_ENABLE_GPU": 1} if with_cuda else {}), ld_flags=["-Xcompiler", "-fopenmp"] if with_omp else [], is_cpp=True, use_cuda_if_available=True, verbose=verbose) tf_mod = compiler.load_tf_module() assert hasattr(tf_mod, "WarpRNNT"), "content of mod: %r" % (dir(tf_mod), ) _tf_mod = tf_mod return tf_mod
def init_warprna(verbose=False): """ Initializes and compiles the library. Caches the TF module. :param bool verbose: """ global _tf_mod if _tf_mod: return assert is_checked_out( ), "submodule not checked out? Run `git submodule update --init --recursive`" assert is_tf_cuda_build(), "TF with CUDA support required" src_files = [ '%s/tensorflow_binding/src/warp_rna_op.cc' % submodule_dir, '%s/core.cu' % submodule_dir, ] src_code = "" for fn in src_files: f_code = open(fn).read() src_code += "\n// ------------ %s : BEGIN { ------------\n" % os.path.basename( fn) # https://gcc.gnu.org/onlinedocs/cpp/Line-Control.html#Line-Control src_code += "#line 1 \"%s\"\n" % os.path.basename(fn) src_code += f_code src_code += "\n// ------------ %s : END } --------------\n\n" % os.path.basename( fn) compiler = OpCodeCompiler(base_name="warprna_kernels", code_version=1, code=src_code, include_paths=(submodule_dir, os.path.dirname(submodule_dir)), c_macro_defines={"WARPRNA_ENABLE_GPU": 1}, is_cpp=True, use_cuda_if_available=True, verbose=verbose) tf_mod = compiler.load_tf_module() assert hasattr(tf_mod, "WarpRNA"), "content of mod: %r" % (dir(tf_mod), ) _tf_mod = tf_mod return tf_mod
def get_filtered_score_op(verbose=False): cpp_code = """ #include "tensorflow/core/framework/op.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/shape_inference.h" #include "tensorflow/core/framework/resource_mgr.h" #include "tensorflow/core/framework/resource_op_kernel.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/platform/macros.h" #include "tensorflow/core/platform/mutex.h" #include "tensorflow/core/platform/types.h" #include "tensorflow/core/public/version.h" #include <cmath> #include <map> #include <set> #include <string> #include <tuple> using namespace tensorflow; REGISTER_OP("GetFilteredScore") .Input("prev_str: string") .Input("scores: float32") //.Input("labels: string") .Output("new_scores: float32") .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { c->set_output(0, c->input(1)); return Status::OK(); }); class GetFilteredScoreOp : public OpKernel { public: using OpKernel::OpKernel; void Compute(OpKernelContext* context) override { const Tensor* prev_str = &context->input(0); const Tensor* scores = &context->input(1); //const Tensor* labels = &context->input(2); int n_batch = prev_str->shape().dim_size(0); int n_beam = prev_str->shape().dim_size(1); Tensor* ret; OP_REQUIRES_OK(context, context->allocate_output(0, TensorShape({n_batch, n_beam}), &ret)); for(int bat = 0; bat < n_batch; ++bat) for(int hyp = 0; hyp < n_beam; ++hyp) ret->tensor<float, 2>()(bat, hyp) = scores->tensor<float, 2>()(bat, hyp); for(int bat = 0; bat < n_batch; ++bat) { std::map<tstring, std::set<int> > new_hyps; // seq -> set of hyp idx for(int hyp = 0; hyp < n_beam; ++hyp) { auto& seq_set = new_hyps[prev_str->tensor<tstring, 2>()(bat, hyp)]; seq_set.insert(hyp); } for(const auto& items : new_hyps) { if(std::get<1>(items).size() > 1) { float best_score = 0.; int best_idx = -1; for(int idx : std::get<1>(items)) { float score = scores->tensor<float, 2>()(bat, idx); if(score > best_score || best_idx == -1) { best_score = score; best_idx = idx; } } float sum_score = 0.; for(int idx : std::get<1>(items)) { float score = scores->tensor<float, 2>()(bat, idx); sum_score += expf(score - best_score); } sum_score = logf(sum_score) + best_score; for(int idx : std::get<1>(items)) { if(idx != best_idx) ret->tensor<float, 2>()(bat, idx) = -std::numeric_limits<float>::infinity(); else ret->tensor<float, 2>()(bat, idx) = sum_score; } } } } } }; REGISTER_KERNEL_BUILDER(Name("GetFilteredScore").Device(DEVICE_CPU), GetFilteredScoreOp); """ from returnn.tf.util.basic import OpCodeCompiler compiler = OpCodeCompiler(base_name="GetFilteredScore", code_version=1, code=cpp_code, is_cpp=True, use_cuda_if_available=False, verbose=verbose) tf_mod = compiler.load_tf_module() return tf_mod.get_filtered_score