How to convert multi-label-classifier to torch script in pytorch 1.0

jit
1

Hi Everyone,

I’m trying to convert multilabel classification model (pytorch 1.0) to torch script via annotation.

Here is my model:

class MultiLabelModel(nn.Module):
    def __init__(self, basemodel, basemodel_output, num_classes):
        super(MultiLabelModel, self).__init__()
        self.basemodel = basemodel
        self.num_classes = num_classes
        for index, num_class in enumerate(num_classes):
            setattr(self, "FullyConnectedLayer_" + str(index), nn.Linear(basemodel_output, num_class))
    
    def forward(self, x):
        x = self.basemodel.forward(x)
        outs = list()
        dir(self)
        for index, num_class in enumerate(self.num_classes):
            fun = eval("self.FullyConnectedLayer_" + str(index))
            out = fun(x)
            outs.append(out)
        return outs

The basemodel is predefined backbone model.

I tried this:

class MultiLabelModel(torch.jit.ScriptModule):
    __constants__ = ['num_classes']
    __constants__ = ['basemodel']
    __constants__ = ['basemodel_output']
    def __init__(self, basemodel, basemodel_output, num_classes):
        super(MultiLabelModel, self).__init__()
        self.basemodel = basemodel
        self.num_classes = num_classes
        self.basemodel_output = basemodel_output
        for index, num_class in enumerate(num_classes):
            setattr(self, "FullyConnectedLayer_" + str(index), nn.Linear(basemodel_output, num_class))
    @torch.jit.script_method
    def forward(self, x):
        x = self.basemodel.forward(x)
        outs = list()
        for index, num_class in enumerate(self.num_classes):
            fun = eval("self.FullyConnectedLayer_" + str(index))
            out = fun(x)
            outs.append(out)
        return outs

and get error:

expected a value of type Tensor for argument ‘0’ but found int

seems like I can’t define the multilabel classifier like this:

        for index, num_class in enumerate(num_classes):
            setattr(self, "FullyConnectedLayer_" + str(index), nn.Linear(basemodel_output, num_class))

Hope somebody can help me out.

2

Could you share a complete example which includes driver code for the script?

3

Complete example:

import torch
import torchvision
import torch.nn as nn

MyClassNum = 10


class FeatureExtraction(torch.nn.Module):
    def __init__(self):
        super(FeatureExtraction, self).__init__()
        self.resnet = torchvision.models.resnet18(pretrained=True)
        self.resnet = nn.Sequential(*list(self.resnet.children())[:-1])
        self.resnet.cuda()
    def forward(self, image_batch):
        return self.resnet(image_batch)

class MultiLabelModel(torch.jit.ScriptModule):
    __constants__ = ['num_classes']
    def __init__(self, num_classes):
        super(MultiLabelModel, self).__init__()
        self.num_classes = MyClassNum
        for index in enumerate(range(0,self.num_classes)):
            setattr(self, "FullyConnectedLayer_" + str(index[0]), nn.Linear(512, 1))
    @torch.jit.script_method
    def forward(self, x):
        x = x.view(x.size(0), -1) # flatten
        outs = list()
        for index in enumerate(range(0,self.num_classes)):
            fun = eval("self.FullyConnectedLayer_" + str(index[0]))
            out = fun(x)
            outs.append(out)
        return outs

class MultiLabelClassifier(torch.jit.ScriptModule):
    def __init__(self, num_classes):
        super(MultiLabelClassifier, self).__init__()
        self.FeatureExtraction = FeatureExtraction()
        self.classifier = MultiLabelModel(num_classes)
    @torch.jit.script_method
    def forward(self, img):
        feature = self.FeatureExtraction(img)
        preLabels= self.classifier(feature)
        return preLabels

my_script_module = MultiLabelClassifier(MyClassNum)
my_script_module.save("model.pt")

Thanks for your quick reply.