Quick start example with GPU

[AZaitzeff]

In your quick start example you have an example for cpu:
# Extract per-residue embeddings (on CPU) with torch.no_grad(): results = model(batch_tokens, repr_layers=[34]) token_embeddings = results["representations"][34]
if I wanted to Extract per-residue embeddings on a gpu how would I change the above lines?

Thanks

[tomsercu]

Hi - thanks for your question!
You can look at extract.py which does embeddings on gpu in batch, specifically look for this kind of logic around transfering to device: https://github.com/facebookresearch/esm/blob/master/extract.py#L62-L64
You can also look at pytorch introductory tutorials to learn how to work with devices in pytorch.
Cheers,
Tom

[joshim5]

In case it helps, here's a quick colab notebook with a concrete example. It's as easy as calling .cuda() on the model and on the data.

import torch

model, alphabet = torch.hub.load("facebookresearch/esm", "esm1_t34_670M_UR50S")
batch_converter = alphabet.get_batch_converter()

# Prepare data (two protein sequences)
data = [("protein1", "MYLYQKIKN"), ("protein2", "MNAKYD")]
batch_labels, batch_strs, batch_tokens = batch_converter(data)

# Extract per-residue embeddings (on GPU)
model = model.cuda()
with torch.no_grad():
    results = model(batch_tokens.cuda(), repr_layers=[34])
token_embeddings = results["representations"][34]

[AZaitzeff]

@joshim5 Thank you! That helps a lot.