How to “lightninfy” the official PyTorch sentiment analysis tutorial?

[davidefiocco]

Hi, I'm trying to refactor the official NLP (sentiment analysis) tutorial, using Lightning in order to take advantage of things like early stopping etc.

I'm moving first steps, and the main hurdle is the creation of a Lightning module, and in particular coding the training_step.

What I came up so far is

class LitTextClassifier(pl.LightningModule):
    def __init__(self, num_class, criterion = CrossEntropyLoss):
        super().__init__()
        self.embedding = nn.EmbeddingBag(VOCAB_SIZE, EMBED_DIM, sparse=False)
        self.fc = nn.Linear(EMBED_DIM, num_class)
        self.init_weights()
        self.criterion = criterion

    def init_weights(self):
        initrange = 0.5
        self.embedding.weight.data.uniform_(-initrange, initrange)
        self.fc.weight.data.uniform_(-initrange, initrange)
        self.fc.bias.data.zero_()

    def forward(self, text, offsets):
        embedded = self.embedding(text, offsets)
        return self.fc(embedded)

    def configure_optimizers(self):
        optimizer = optim.SGD(self.parameters(), lr=4.0)
        return optimizer

    def training_step(self, batch, batch_idx):
        # I am messing up things here
        text, offsets, cls = batch
        output = self.forward(text, offsets)
        loss = self.criterion(output, cls)

        return loss

I think I am getting the training_step wrong. Can someone provide guidance here?
A full gist to reproduce code + errors I get is here: https://gist.github.com/davidefiocco/3b6c6b1e09c4f664b3a73e5bf24d1668/5aa4c224f7772db835bbaa92d559837c7a40f4df

[akihironitta]

@davidefiocco Hi, I think you're trying to instantiate the criterion class with output and cls. You need to instantiate it in advance:

-        self.criterion = criterion
+        self.criterion = criterion()

[davidefiocco]

Indeed @akihironitta !! Silly me... following your advice https://gist.github.com/davidefiocco/3b6c6b1e09c4f664b3a73e5bf24d1668/14eb90a1ec5ecfd1f7e92a7c3b9b1e554ccc81fa should now be a better version!