[code_search] Train a high quality model

[jlewi]

We'd like to train a high quality model for the code search example.

I took an initial stab at this.
My parameters are here:https://github.com/jlewi/examples/blob/7346d6f1605eaf3e9251e3eec7b1802f5683e637/code_search/kubeflow/components/params.libsonnet#L33

• I ran for ~300K steps
• I used 1 K80 GPU
• Eval steps = 10000
• Used the transformer_tiny hyperparameters set.

Results from TensorBoard are below

Accuracy and eval appear to flatten out after about 150K steps.

@activatedgeek Any suggestions about what experiments to run?

[activatedgeek]

These metrics are probably not the right thing to check straight away. Here are the reasons - 1. The model accuracy is not exactly the classification accuracy because we are using the special keys "targets" to mean the final label (which in our case is a string embedding). I think the internal T2T code is not considering this as a classification problem. This is why those accuracy values are abysmally low (Order of micro). 2. The error currently calculated also includes the padding tokens so isn't exactly accurate. Fixing this will certainly cause loss to go further down. Let me look at the codebase and try to figure out how to get T2T using the right inputs. But I am pretty sure the events recorded here are not reporting what they seem to be reporting.

…

On Mon, Aug 27, 2018 at 7:51 PM Jeremy Lewi ***@***.***> wrote: We'd like to train a high quality model for the code search example. I took an initial stab at this. My parameters are here: https://github.com/jlewi/examples/blob/7346d6f1605eaf3e9251e3eec7b1802f5683e637/code_search/kubeflow/components/params.libsonnet#L33 - I ran for ~300K steps - I used 1 K80 GPU - Eval steps = 10000 - Used the transformer_tiny hyperparameters set. Results from TensorBoard are below Accuracy and eval appear to flatten out after about 150K steps. @activatedgeek <https://github.com/activatedgeek> Any suggestions about what experiments to run? [image: loss_eval] <https://user-images.githubusercontent.com/777219/44692792-6d552100-aa19-11e8-93f9-2ce941c8a335.png> [image: accuracy] <https://user-images.githubusercontent.com/777219/44692801-747c2f00-aa19-11e8-99c6-bf89f125dabb.png> [image: loss] <https://user-images.githubusercontent.com/777219/44692809-7940e300-aa19-11e8-860a-192c697df7bc.png> — You are receiving this because you were mentioned. Reply to this email directly, view it on GitHub <#239>, or mute the thread <https://github.com/notifications/unsubscribe-auth/ADupLUavb1i7ECLSknPtdXBM_rw82JQ7ks5uVIYfgaJpZM4WOroL> .

-- Regards, Sanyam Kapoor https://blog.sanyamkapoor.com

[jlewi]

Thanks.

@cwbeitel Any additional ideas?

[cwbeitel]

Sure. Just so I'm clear, correct me if this is wrong, but the idea here is to learn a mapping between a doc string and the code it corresponds to, then perform search by mapping a new query into embedding space and returning code snippets with embeddings in a close neighborhood of the query? So your loss above corresponds to the first stage, where we're learning a mapping between doc strings and functions? Are you seeing poor model performance when attempting to re-purpose that embedding for search or poor performance on the original task as well? Can your provide some example code that is produced for some queries that are in and not in your dataset?

Also what do these orange and blue lines correspond to? Training and eval? If that's your training loss curve then something is off with the problem setup. What loss function are we using?

You might try constructing a much simpler toy problem that is similar to this one to do some basic debugging. For example only work with doc string / code pairs that are really short (both short doc string and small amounts of code). Then you can progressively increase the length considered with training on shorter snippets serving as pre-training for longer snippets. You can imagine the problem of generating long blocks of code is kind of a challenge but it might be easier if you were first able to generate short ones.

[jlewi]

Thanks @cwbeitel . The network output is the similarity (cosine distance?) between a doc string and code. We have positive examples (want the distance to be small) and negative examples (want the distance to be high).

blue line - is eval (according to label in tensorboard)
red line - doesn't say (I assume its training)

I don't think we've gotten to the point of evaluating qualitatively how well the model is performing on the search task.

I think at this point we simple want metrics that can be used to detect convergence during training and for hyperparameter tuning.

@activatedgeek regarding your first point. Would it be better to use a classification loss metric (e.g. +1, -1) or a regression loss?

/cc @hamelsmu @activatedgeek

[cwbeitel]

Ah that's great, reminds me of this paper from Arandjelovic and Zimmerman (2017) that trained a model to predict the correspondence of audio and visuals in a corpus of videos.

Looks like the batch size for transformer_tiny that is inherited from transformer_base_v1 is 4096 so it's probably not too small but I'd wonder if maybe your examples aren't shuffled properly. If that was the case (or your batch size was too small) you might see instability in the loss fn like you're showing above as a result of training batches not being representative of the full distribution.

[activatedgeek]

@jlewi The loss metric might need improvement but I think we can get mileage without making any changes to that. I suspect it is because the "target" (from the "features" dict) in our case is a continuous high dimensional vector. When compared against the model output (during eval), it is leading to almost zero accuracy because using the accuracy metric for continuous domain wouldn't make sense.

@cwbeitel The data shuffling is done by T2T itself but that could certainly be one place to play with.

[jlewi]

@activatedgeek Thanks. Can you suggest:

1. what experiments we should run to train a model?
2. what metrics we should look at to pick the best model?
3. suggestion about qualitatively evaluating the final product/model?

I tried using hparams transformer_base_single_gpu
https://github.com/jlewi/examples/blob/jlewi_train/code_search/kubeflow/components/params.libsonnet#L47

Here's a graph of the loss. eval loss seems to flatten out after only 4K trials.

My job failed after ~120K steps and I'm not quite sure why. There was no error in stdout/stderr of the process so I expect it exited with some non-zero exit code for some reason. kubeflow/training-operator#811 will hopefully provide better logging.

[cwbeitel]

Random thought - you could restrict your training to a single language or a query subset of that language to make the learning problem simpler, either for debug or towards using a search approach that ensembles many specialized models.

[activatedgeek]

@jlewi Sorry about the delay. Moving and a lot of first week formalities.

I am trying to fix the loss function computation. Once I see new graphs with that, I'll revert back.

@cwbeitel That is a great point. I am trying to overfit the model to a very small dataset to make sure the training is working as intended.

[jlewi]

I am trying to fix the loss function computation. Once I see new graphs with that, I'll revert back.

Typo? Did you mean you will create a PR with the new loss function?

[cwbeitel]

@activatedgeek It looks like the model Hamel trained in Part 4 of his series on code search (https://github.com/hamelsmu/code_search/blob/master/notebooks/4%20-%20Train%20Model%20To%20Map%20Code%20Embeddings%20to%20Language%20Embeddings.ipynb) started with a pre-trained language model on doc strings (described in Part 3). Do you think doing something similar here (as well as perhaps replicating his choice of model) would be helpful?

[activatedgeek]

@cwbeitel We did consider that at the start however went in favor of training an end-to-end task specific model instead of doing a two step thing - first and language model and then fine-tuning for the task.

[cwbeitel]

@activatedgeek Because that makes continuous training more straightforward or something else? Not to knock the approach, it seems like it should work.

[activatedgeek]

@cwbeitel Certainly, only to make it an end-to-end architecture and less moving parts to deal with.

[activatedgeek]

I just made some new changes in #233 but after all the training it crashes up with dimension errors.

[hamelsmu]

@inc0 maybe you can take a look today after we talk.

[cwbeitel]

Hmm sorry to hear that, can you roll back to what was working and see which of the additions in the current version cause that problem?

Also debugging interactively with tf.Eager (e.g. https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/notebooks/hello_t2t.ipynb) or some local tests (e.g. https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/models/image_transformer_2d_test.py) could be helpful.

[activatedgeek]

@cwbeitel The code on the master works just that during inference it always ends up using one of the networks. I wanted to fix that by adding a TensorFlow conditional but broke training.

[cwbeitel]

Maybe it would be helpful to have a test with tiny hparams that runs training and inference in series locally. In my opinion tf.Eager makes this more straightforward, you could just fork this one: https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/data_generators/allen_brain_test.py#L157

[activatedgeek]

@cwbeitel Do you know of an easy way to convert my non-eager T2T code to eager mode? This would be a great savior to help me debug.

[cwbeitel]

The model will work with both. In that test model = image_transformer_2d.Img2imgTransformer(...) instantiates a typical model built in the same way as yours, decorates a loss function to work with eager, then directly applies gradients produced from it. Having called tfe.enable_eager_execution() on startup.

[activatedgeek]

@cwbeitel Do you mind skimming through the two files (function_docstring.py and similarity_transformer.py) in #233? I'm kind of lost in the T2T abstractions here.

The error at the end of training is

InvalidArgumentError (see above for traceback): logits and labels must be broadcastable: logits_size=[6,128] labels_size=[84,128]
	 [[{{node cs_similarity_transformer/parallel_0_5/cs_similarity_transformer/cs_similarity_transformer/padded_cross_entropy/smoothing_cross_entropy/softmax_cross_entropy_with_logits}} = SoftmaxCrossEntropyWithLogits[T=DT_FLOAT, _device="/job:localhost/replica:0/task:0/device:CPU:0"](cs_similarity_transformer/parallel_0_5/cs_similarity_transformer/cs_similarity_transformer/body/cond/Merge, cs_similarity_transformer/parallel_0_5/cs_similarity_transformer/cs_similarity_transformer/padded_cross_entropy/smoothing_cross_entropy/one_hot)]]

When I put --eval_steps=1, this error goes away.

[inc0]

What I've noticed (yesterday evening, didn't get too much into it) is that Embedding layers in T2T are much simpler than what Hamel did (or it looks like it to me).
Particularly, if I understand correctly, this:
https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/models/transformer.py#L1081 seems to be correspondent with this:
https://github.com/hamelsmu/code_search/blob/master/notebooks/seq2seq_utils.py#L48-L80
So we're replacing encoder/decoder+embeddings with just embeddings tested for english nlp. Not sure how much does that affect accuracy. Another thing is same function in t2t is used for language modelling, where in Hamel's example it's dedicated FastAI model (which I haven't dig into yet). After that we're simply calculating distances with t2t. In effect T2T models seems quite a bit shallower than what Hamel initially created, maybe that's affecting accuracy.

Let me dig a little more into this, T2T isn't exactly best documented project so I'm shooting blind and missing stuff:)

[cwbeitel]

Ok hm well just to confirm regarding the shuffling question I raised a while ago it looks like for Text2TextProblem this is handled in for you but you could also confirm that by iterating over examples in a notebook and Googling the function docstrings for sequential examples.

Also more generally it would be good to just look at what is ending up in "inputs", "targets", and "emb_code" just to sanity check.

Looking at similarity_transformer.py and based on another discussion what comes to mind is whether the variables for the two encoder branches are being initialized in the right way. Idk if in tf.variable_scope if the initializer isn't specified it just inherits from the parent scope and if none is specified it just uses zero? You could try dropping in a specific initializer to those or maybe I'm misunderstanding how that works 😜

The error you're getting looks like you're trying to compute the cross entropy of a prediction with a different shape from a ground-truth example. Can you roll back to a state of the code where this didn't happen and step forward to see what change causes this? Doing this locally against a tf.Eager test that tests both training and inference should save you a lot of time relative to launching new runs on Kubeflow if that's what you're doing currently.

Also what is your reasoning around this tf.cond and the use of 'embed_code'? It looks like it's always zero or am I missing something?

Also FYI you can directly examine the values for the gradients of each layer at each iteration by examining gv in loss, gv = loss_fn(example), in the context of the tf.Eager thing I referenced above.

[cwbeitel]

And you said that error happens at the end of training, not during? Then it's probably a problem with the eval phase specifically. If you turn off eval by using schedule=train then you shouldn't see the error. i.e. toward identifying whether that's where the problem is arising.

[cwbeitel]

And you probably will want a way to sanity check the cosine distance code in body as @jlewi as suggesting.

[jlewi]

Filed #254 about measuring the quality of the model since this issue was getting logged.

I think once we fix #254 we could do a quick sanity check of comparing a trained model to a randomly initialized model.

[cwbeitel]

@activatedgeek Maybe this notebook will make debugging a little easier!

[jlewi]

Thanks @cwbeitel

@cwbeitel @activatedgeek @hamelsmu It looks like the vocabulary in
gs://kubeflow-examples/t2t-code-search/20180802/data/vocab.github_function_docstring.8192.subwords

Has 8164 words. Does that seem reasonable?

[jlewi]

Just noticed @cwbeitel's comment #259 (comment) in the other issue noting that the vocab size seems low.

[cwbeitel]

You could get a sense of it by taking a sample of raw doc string or code and compute edit_distance(decode(encode(string)), string) where decode and encode are coders that use that vocab, maybe more or less this:

# Get the encoders from the problem
encoders = problem_object.feature_encoders(data_dir)

rt = encoders["inputs"].decode(encoders["inputs"].encode(input_str))

d = edit_distance(rt, input_str)

[cwbeitel]

This results from the setting of approx_vocab_size (where 2^13 ~ 8k), e.g.

@registry.register_problem
class GithubFunctionDocstring(text_problems.Text2TextProblem):

  @property
  def approx_vocab_size(self):
    return 2**13
...

Numerous other language-oriented problems in t2t use vocab sizes of 8k or 32k, see this... but perhaps a larger vocab is needed for code than prose given the ways words combine with symbols, e.g. here's a sample from my vocab:

'<pad>'
'<EOS>'
'\u_'
'self_'
'_'
's_'
'if_'
'def_'
'the_'
'return_'
'0_'
'in_'
'1_'
'for_'
'a_'
' \u_'
'name_'
'to_'
' ._'
'get_'
'is_'
'n_'
'None_'
...
'figs'
'exi'
'ency'
'contrib'
'chat_'
'bject'
'area'
'arc_'
'anima'
'agg_'
'addCa'
'Windows_'
'Mus'
'Left'
'Forma'
'Channel'
'0001_'
'{'
'}'
'|'
'~'
'?'
'&'
'.'
'<'
'$'
'@'
'^H'
'^G'
'ü'
'ó'
'³'
'^P'
'^D'
'^B'
'^A'

Also it's worth noting that the language translation problems I've seen in the t2t repo have vocabs around 8-32k per language and the code and doc strings we're working with are both drawn from a variety of languages, whether spoken or programmatic.

The easy route is to just try increasing the vocab and see if it improves the quality. Or we might see what happens when we filter the dataset e.g. to be only english python (or perhaps further to be that with nicely formatted doc strings e.g. including an Args: section).

[jlewi]

Thanks @cwbeitel

[stale]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.