How to train custom Word Embeddings using GPU on AWS | by Sanket Gupta | Towards Data Science

Language is important. Humans use words to communicate, and they carry meaning. Can we train a machine to also learn meaning? The topic of this blog is how we can train a machine to learn meaning of words using word embeddings. As I go through the loop of Stack Overflow, EC2 documentations and blogs, I am going to jot down the steps for someone to train word embeddings using TensorFlow on a GPU on AWS. But before going through the steps, let’s briefly understand why we need word embeddings in the first place.
Typically, in most natural language processing tasks words are converted to one-hot vectors where each word is gets a unique index in a vector. So in a 3 word vocabulary of apple, orange and juice, we can make [1 0 0], [0 1 0] and [0 0 1] to represent these 3 words. For small language tasks, these vectors are quite useful. However, there are 2 main disadvantages of these: a) For a training data with n words, we will need n dimensional vectors. So you can imagine 1,000,000 dimensional vectors which are quite large and will be quite sparse. b) There is no sense of context in one-hot vectors — each word is completely unique. So if a model has learnt “apple juice”, it can’t transfer the learning to fill “orange ___ “. Word embeddings follow the principle “you shall a word by the company it keeps”, and so we each dimension in a word embedding vector could represent a “feature” such as food or gender or capital and so on — as can be seen below:
Well trained word embeddings can be used for translation, text summarization, sentiment analysis, and analogy tasks such as man > woman so king > queen. Word2Vec and Glove are two popular pre-trained word embeddings available as open-source — they are quite powerful to get started with general language tasks. But, in this tutorial we will learn how to build the word embeddings for our own custom data in a particular niche. I am working to build word embeddings for electronics data such as microprocessors, capacitors, RAMs, power supplies etc.
So, let’s get started on the journey to make word embeddings which we’ll train on an AWS GPU!
Choose P2 or P3 instance typeFirstly, you need to choose an EC2 instance type. Recently AWS released p3 instances on EC2, but they cost ~$3 an hour which is quite a lot. So I am going to stick with p2 instance which might not be as fast but cheaper (around $1 per hour). This blog goes in more detail for a comparison. Once you have chosen p2 instance, you might need to request a limit increase on your account to get to use them.Set up a Deep Learning AMI and Security GroupsIf you were building up your own GPU hardware, you’d need to set up CUDA drivers on your own but using EC2 instances on AWS gives us the benefit of using pre-built AMI or Amazon Machine Image. I chose Deep Learning AMI with Source Code (CUDA 8, Ubuntu). UPDATE (June 2018): Latest TensorFlow libraries are not quite compatible with CUDA 8, please use an AMI which provides support for CUDA 9 such as this one.
Choose P2 or P3 instance typeFirstly, you need to choose an EC2 instance type. Recently AWS released p3 instances on EC2, but they cost ~$3 an hour which is quite a lot. So I am going to stick with p2 instance which might not be as fast but cheaper (around $1 per hour). This blog goes in more detail for a comparison. Once you have chosen p2 instance, you might need to request a limit increase on your account to get to use them.
Set up a Deep Learning AMI and Security GroupsIf you were building up your own GPU hardware, you’d need to set up CUDA drivers on your own but using EC2 instances on AWS gives us the benefit of using pre-built AMI or Amazon Machine Image. I chose Deep Learning AMI with Source Code (CUDA 8, Ubuntu). UPDATE (June 2018): Latest TensorFlow libraries are not quite compatible with CUDA 8, please use an AMI which provides support for CUDA 9 such as this one.
If we want to use Jupyter notebook in our EC2 instance, we need to change the security groups to allow access to the EC2 from port 8888. This can be done by creating a new security group and using Custom TCP Rule and Port Range as 8888 and in source select Anywhere. You can now launch the instance. Next, create a new key pair and save it locally on your machine. I have had to do these steps again when I lost this .pem file so make sure its in a safe place! Change the read-write permissions to readonly using chmod 400 YourPEMFile.pem. Next, go ahead and launch your GPU by right click on the instance Instance State > Start.
3. Login and set up Conda environment on EC2 instanceYou can login to your GPU by ssh -i path/YourPEMFile.pem ubuntu@X.X.X.X where X.X.X.Xrepresents the Public IP of your GPU and path/YourPEMFile.pemis the path of your .pemfile you created in the previous step. Once you are logged into the EC2 instance, you can set up your own Conda environment. Why do you need it? Well read this answer which discusses how Conda is better than virtualenv. Typing conda -V will result in an error as the PATH variables havn’t been set up yet. We need to go to .bash_profile and add the PATH.
Here are the steps to get the conda working on the new EC2 instancea. Create or edit the file byvim ~/.bash_profile and b. Add export PATH=”/home/ubuntu/src/anaconda2/bin:$PATH” PATH specifies the directories which have executable programs that can be started without knowing the full path on the command line.c. Run source ~/.bash_profile to source this file. d. You can now run conda -V and will not see an error.e. Create a new environment by conda create -n embeddings python=3 This will create a new conda environment with latest version of Python 3. f. Run source activate embeddings to get the conda environment started.
4. Install Tensorflow-GPU, Jupyter notebook and other packagesHaving set up and activated the conda environment, you can now install the Tensorflow and other packages. a) We can usepip install tensorflow-gpu jupyter notebook nltk scikit-learn pandas to install all the necessary packages to do the word embeddings stuff. We need to use tensorflow-gpuand not tensorflowas the former is optimized to be used in GPU. All these requirements can be saved into a file using pip freeze > requirements.txt
b) We can now start the jupyter notebook using jupyter notebook --ip=0.0.0.0 --no-browser Go to X.X.X.X:8888/?token=<TOKEN> where X.X.X.X is the IP address of your EC2 machine and <TOKEN> is the token which comes up when we run the command above.
c) Run following code to check TensorFlow version and see that the TensorFlow is correctly running on GPU.
import tensorflow as tfprint('TensorFlow Version: {}'.format(tf.__version__)) print('Default GPU Device: {}'.format(tf.test.gpu_device_name()))
The output I get is:TensorFlow Version: 1.3.0Default GPU Device: /gpu:0
To be absolutely sure that GPU is actually being used, runsess = tf.Session(config=tf.ConfigProto(log_device_placement=True)) and you should see something ending with /gpu:0 which is also suggested by official documentation. This helped me recently in setting up to actually use the GPU.
5. Build corpus of text to be used for training the word embeddingsYou could use IMDB or Wikipedia dataset, but since the goal of this blog is to work through custom word embeddings, we are going to build our own corpus. If you have a lot of text files, then you can read them one by one and build the corpus. Here, I am assuming you have PDF files which are little harder to work with. I am using the library PyPDF2 for this purpose. The code I used is:
import PyPDF2def read_pdf(fname):    pdfFileObj = open(fname, 'rb')    pdfReader = PyPDF2.PdfFileReader(pdfFileObj)    print("TOTAL PAGES in file", pdfReader.numPages)    pageObj = pdfReader.getPage(0)    text = pageObj.extractText()    return ' '.join(text.split())
I ran the above code over a bunch of PDF files stored locally. I concatenated the strings returning from the above function and wrote into a text file. Our corpus of text is ready!
6. Preprocess the text and create lookup tablesNow, let’s preprocess the text to get it ready for training the word embeddings. In Natural Language Processing, before running a classifier or an algorithm, it’s a common practice to lower case all the words, and do some stop word removal. We are going to use the library re (regular expression library in Python) to perform some of these tasks.
text = text.lower()text = text.replace('.', ' __PERIOD__ ') text = text.replace('-', ' ')    # Remove ' and full stops and bracketstext = re.sub(r'[{}()\']', '', text)# Replace commas, dashes, colons and @ with spacestext = re.sub(r'[;,:-@#]', ' ', text) 
We can do more pre-processing as required, like changing abbreviations to full forms, replacing numbers with __NUMBER__ token and so on. Above, I am changing the full stop to a __PERIOD__ which can be used as a token to train on. It will be important to understand the dynamics of word context in a sentence, and __PERIOD__ will help us do that.
Next, let’s create a couple of hash tables which we can use for lookup of vocabulary both from index > vocabulary as well as vocabulary > index. The following code does this:
word_counts = Counter(words)sorted_vocab = sorted(word_counts, key=word_counts.get, reverse=True)int_to_vocab = {ii: word for ii, word in enumerate(sorted_vocab)}vocab_to_int = {word: ii for ii, word in int_to_vocab.items()}
7. Perform subsamplingWords that occur often such as “the”, “is” etc are not very useful for providing context to nearby words. If we remove all of them, we are effectively removing any information they provide. A better approach as proposed by Mikolov et al. is to remove some of these words to remove some of the noise from our data. For each word in training set, we discard it with probability:
where t is a threshold parameter and f(w) is the frequency of the word w.
8. Create TensorFlow Graph
I always had a hard time visualizing word embeddings and how the math actually works. I found Andrew Ng’s explanation quite useful in understanding the math, and here is a workflow of what we are trying to do with TensorFlow (assuming a word vector with 300 dimensions and a vocabulary of 100,000 unique words):
The input words are passed as one-hot vectors which go into a hidden layer of linear units. These are then connected to a soft-max layer which is used to predict the context words. In other words, given every word, we will try to minimize the loss in predicting neighbor or context words. We can just throw away the soft-max layer weights in the end, and just use the embedding matrix of the word vectors that we have built — these word vectors are lower dimensional representation of the words while preserving the context! Chris McCormick’s blog was an inspiration for the following diagram which is another view of the graph above:
TensorFlow uses dataflow graphs to compute deep learning operations, and we need to first define the graph and then run our data on this graph in a session. Let’s run through this code:
import tensorflow as tftrain_graph = tf.Graph()  # Placeholders for our training data with train_graph.as_default():    inputs = tf.placeholder(tf.int32, [None], name='inputs')    labels = tf.placeholder(tf.int32, [None, None], name='labels')n_vocab = len(int_to_vocab)n_embedding = 300  # Number of embedding features# Embedding matrix that we care about with train_graph.as_default():    embedding = tf.Variable(tf.random_uniform((n_vocab, n_embedding), -1, 1))    word_vector = tf.nn.embedding_lookup(embedding, inputs)
First we are defining inputs and labels which are placeholders in the graph. Next, we are defining what we want to train on using Variable . This discussion on their difference is very good. We use a feed_dict to feed values into a placeholder. Our total vocabulary size x number of embedding dimensions gives us the total number of weights we need to train. Next, we can define the softmax weights and biases:
# Softmax layer that we use for training but don't care aboutwith train_graph.as_default():    softmax_w = tf.Variable(tf.truncated_normal((n_vocab, n_embedding), stddev=0.1))    softmax_b = tf.Variable(tf.zeros(n_vocab)) 
9. Negative Sampling and Loss FunctionWe have prepared a set of words and context words which are positive examples. But we will also add a few words and non-context words as well. These will serve as negative examples. So orange juice is positive example, but also orange book and orange king are negative examples. This makes the learning of the vectors better. For every positive example we can choose between 2 to 20 negative examples. This is called negative sampling and TensorFlow has a function tf.nn.sampled_softmax_loss for this!
We already have written up the entire graph, we need to write what’s the loss function and how we are going to optimize the weights.
with train_graph.as_default():   loss = tf.nn.sampled_softmax_loss(softmax_w, softmax_b,                                       labels, word_vector,                                      n_sampled, n_vocab)       cost = tf.reduce_mean(loss)   optimizer = tf.train.AdamOptimizer().minimize(cost)
10. Run TensorFlow Session to train word embeddingsHaving written out the graph and the loss function as well as optimizer, we can write the TensorFlow session to train the word embeddings.
with tf.Session(graph=train_graph) as sess:    loss = 0    sess.run(tf.global_variables_initializer())    for e in range(1, 11): # Run for 10 epochs        batches = get_batches(train_words, batch_size, window_size)        for x, y in batches:                        train_loss, _ = sess.run([cost, optimizer], feed_dict= {inputs: x, labels: np.array(y)[:, None]})            loss += train_loss
get_batches function is not shown here, but is in the notebook (in references) — it returns a generator of features and classes in batches.The above step can take a long time on a CPU but runs quite fast on the p2.x instance. I was able to run the 10 epochs in a few minutes for training word size of 100k words. We can use the embedding we have learnt to plot T-SNE graph to visualize how the contexts have been learnt. Note that T-SNE graphs is going to reduce our 300D data into 2D so some of the context is going to be lost, however it retains a lot of the information on context.
from sklearn.manifold import TSNEtsne = TSNE() embed_tsne = tsne.fit_transform(embedding[:viz_words, :]) fig, ax = plt.subplots(figsize=(14, 14)) for idx in range(viz_words):          plt.scatter(*embed_tsne[idx, :], color='steelblue')                 plt.annotate(int_to_vocab[idx], (embed_tsne[idx, 0], embed_tsne[idx, 1]), alpha=0.7)
We could also use PCA to visualize the relationships in words, and we would expect to find relationships like these:
This makes word embeddings quite fascinating! If you are working in an e-commerce company, it might make a huge difference to users if you can start identifying if battery size and battery capacity is the same thing, or if you are at a hotel reservation site, it might make a big difference in finding hotel and motel as similar! So word embeddings can be quite powerful in improving experience of users where language is involved.
I find this topic quite interesting too. I love writing and language as well as math and machine learning. Word vectors somehow are the bridge between these two worlds. I just wrapped a large project on text classification, I am going to be working more on these word vector related ideas next!
Good luck~
P.S. I am creating a NEW COURSE on similar topics titled “Overview of Query Understanding Techniques”. This course will cover information retrieval techniques, text mining, query understanding tips and tricks, ways to increase precision, recall and user clicks. Please sign up here if you are interested! This show of interest will allow me to prioritize and build the course.
If you have any questions, drop me a note at my LinkedIn profile or send me an email at sanket@omnilence.com. Thanks for reading!
REFERENCES:
a) Andrew Ng’s course on word embeddings.b) Prof. Manning’s lectures from CS224n at Stanford University. c) Udacity’s Deep Learning Nanodegree.d) My word vector notebook on Github.