I am attempting to use a new NLP model within the PyTorch android demo app Demo App Git however I am struggling to serialize the model so that it works with Android.
The demonstration given by PyTorch is as follows for a Resnet model:
model = torchvision.models.resnet18(pretrained=True)
model.eval()
example = torch.rand(1, 3, 224, 224)
traced_script_module = torch.jit.trace(model, example)
traced_script_module.save("app/src/main/assets/model.pt")
However I am not sure what to use for the 'example' input with my NLP model.
The model that I am using from a fastai tutorial and the python is linked here: model
Here is the Python used to create my model (using the Fastai library). It is the same as in the model link above, but in a simplified form.
from fastai.text import *
path = untar_data('http://files.fast.ai/data/examples/imdb_sample')
path.ls()
#: [PosixPath('/storage/imdb_sample/texts.csv')]
data_lm = TextDataBunch.from_csv(path, 'texts.csv')
data = (TextList.from_csv(path, 'texts.csv', cols='text')
.split_from_df(col=2)
.label_from_df(cols=0)
.databunch())
bs=48
path = untar_data('https://s3.amazonaws.com/fast-ai-nlp/imdb')
data_lm = (TextList.from_folder(path)
.filter_by_folder(include=['train', 'test', 'unsup'])
.split_by_rand_pct(0.1)
.label_for_lm()
.databunch(bs=bs))
learn = language_model_learner(data_lm, AWD_LSTM, drop_mult=0.3)
learn.fit_one_cycle(1, 1e-2, moms=(0.8,0.7))
learn.unfreeze()
learn.fit_one_cycle(10, 1e-3, moms=(0.8,0.7))
learn.save_encoder('fine_tuned_enc')
path = untar_data('https://s3.amazonaws.com/fast-ai-nlp/imdb')
data_clas = (TextList.from_folder(path, vocab=data_lm.vocab)
.split_by_folder(valid='test')
.label_from_folder(classes=['neg', 'pos'])
.databunch(bs=bs))
learn = text_classifier_learner(data_clas, AWD_LSTM, drop_mult=0.5)
learn.load_encoder('fine_tuned_enc')
learn.fit_one_cycle(1, 2e-2, moms=(0.8,0.7))
learn.freeze_to(-2)
learn.fit_one_cycle(1, slice(1e-2/(2.6**4),1e-2), moms=(0.8,0.7))
learn.freeze_to(-3)
learn.fit_one_cycle(1, slice(5e-3/(2.6**4),5e-3), moms=(0.8,0.7))
learn.unfreeze()
learn.fit_one_cycle(2, slice(1e-3/(2.6**4),1e-3), moms=(0.8,0.7))
I worked out how to do this after a while. The issue was that the Fastai model wasn't tracing correctly no matter what shape of input I was using.
In the end, I used another text classification model and got it to work. I wrote a tutorial about how I did it, in case it can help anyone else.
Begin by opening a new Jupyter Python Notebook using your preferred cloud machine provider (I use Paperspace).
Next, copy and run the code in the PyTorch Text Classification tutorial. But replace the line…
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
With…
device = torch.device("cpu")
NOTE: It caused issues tracing when the device was set to CUDA so I forced it on to the CPU. (this will slow training, but inference on the mobile will run at the same speed as it is cpu anyway)
Lastly, run the code below to correctly trace the model to allow it to be run on Android:
data = DataLoader(test_dataset, batch_size=1, collate_fn=generate_batch)
for text, offsets, cls in data:
text, offsets, cls = text.to(device), offsets.to(device), cls.to(device)
example = text, offsets
traced_script_module = torch.jit.trace(model, example)
traced_script_module.save("model.pt")
In addition, if you would like a CSV copy of the vocab list for use on Android when you are making predictions, run the following code afterwards:
import pandas as pd
vocab = train_dataset.get_vocab()
df = pd.DataFrame.from_dict(vocab.stoi, orient='index', columns=['token'])
df[:30]
df.to_csv('out.csv')
This model should work fine on Android using the PyTorch API.