Can anyone provide GPU version of translate.py?

[jonghkim]

Current CPU version of translate.py is so slow. Can anyone provide GPU version of translate.py?

[jonghkim]

I firstly changed multiprocessing version of it to single process version.

export THEANO_FLAGS=device=cuda0,floatX=float32,gpuarray.preallocate=0.9
export CUDA_ROOT=/usr/local/cuda-9.0
export CUDA_HOME=/usr/local/cuda-9.0

export PATH=${CUDA_HOME}/bin:$PATH
export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH

'''
Translates a source file using a translation model.
'''
import argparse

import numpy
import cPickle as pkl

from nmt import (build_sampler, gen_sample, load_params,
                 init_params, init_tparams)

def translate_model(model, options, seqs, k, normalize):

    from theano.sandbox.rng_mrg import MRG_RandomStreams as RandomStreams
    from theano import shared
    trng = RandomStreams(1234)
    use_noise = shared(numpy.float32(0.))
    
    # allocate model parameters
    params = init_params(options)

    # load model parameters and set theano shared variables
    params = load_params(model, params)
    tparams = init_tparams(params)

    # word index
    f_init, f_next = build_sampler(tparams, options, trng, use_noise)

    def _translate(seq):
        # sample given an input sequence and obtain scores
        sample, score = gen_sample(tparams, f_init, f_next,
                                   numpy.array(seq).reshape([len(seq), 1]),numpy.array([1],dtype=numpy.int32),
                                   options, trng=trng, k=k, maxlen=200,
                                   stochastic=False)

        # normalize scores according to sequence lengths
        if normalize:
            lengths = numpy.array([len(s) for s in sample])
            score = score / lengths
        sidx = numpy.argmin(score)
        return sample[sidx]

    traslated = []

    for idx, x in enumerate(seqs):
        seq = _translate(x)
        traslated.append(seq)
        if numpy.mod(idx, 10) == 0:
            print 'Sample ', (idx+1), '/', len(seqs), ' Done'
    
    return traslated


def main(model, dictionary, dictionary_target, source_file, saveto, k=5,
         normalize=False, chr_level=False):
    """
    load model
    """
    # load model model_options
    with open('%s.pkl' % model, 'rb') as f:
        options = pkl.load(f)

    """
    load word embedding info
    """
    # load source dictionary and invert
    with open(dictionary, 'rb') as f:
        word_dict = pkl.load(f)
    word_idict = dict()
    for kk, vv in word_dict.iteritems():
        word_idict[vv] = kk
    word_idict[0] = '<eos>'
    word_idict[1] = 'UNK'

    # load target dictionary and invert
    with open(dictionary_target, 'rb') as f:
        word_dict_trg = pkl.load(f)
    word_idict_trg = dict()
    for kk, vv in word_dict_trg.iteritems():
        word_idict_trg[vv] = kk
    word_idict_trg[0] = '<eos>'
    word_idict_trg[1] = 'UNK'

    # utility function
    def _seqs2words(caps):
        capsw = []
        for cc in caps:
            ww = []
            for w in cc:
                if w == 0:
                    break
                ww.append(word_idict_trg[w])
            capsw.append(' '.join(ww))
        return capsw

    def _get_data(fname):
        with open(fname, 'r') as f:

            seqs = []

            for line in f:
                if chr_level:
                    words = list(line.decode('utf-8').strip())
                else:
                    words = line.strip().split()
                x = map(lambda w: word_dict[w] if w in word_dict else 1, words)
                x = map(lambda ii: ii if ii < options['n_words'] else 1, x)
                x += [0]
                seqs.append(x)

        return seqs

    def _retrieve_seqs(translated_seqs):
        trans = []
        for idx, seq in enumerate(translated_seqs):
            trans.append(seq)
            if numpy.mod(idx, 10) == 0:
                print 'Sample ', (idx+1), '/', len(translated_seqs), ' Done'
        return trans

    print 'Translating ', source_file, '...'
    seqs = _get_data(source_file)
    print 'Get Data Done ...'
    translated_seqs = translate_model( model, options, seqs, k, normalize)
    print 'Translation Done ...'    
    trans = _seqs2words(_retrieve_seqs(translated_seqs))
    print 'Seqs 2 Words Done ...'    
    with open(saveto, 'w') as f:
        print >>f, '\n'.join(trans)
    print 'Done'


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('-k', type=int, default=1)
    parser.add_argument('-n', action="store_true", default=False)
    parser.add_argument('-c', action="store_true", default=False)
    parser.add_argument('model', type=str)
    parser.add_argument('dictionary', type=str)
    parser.add_argument('dictionary_target', type=str)
    parser.add_argument('source', type=str)
    parser.add_argument('saveto', type=str)

    args = parser.parse_args()

    main(args.model, args.dictionary, args.dictionary_target, args.source,
         args.saveto, k=args.k, normalize=args.n, chr_level=args.c)