Converting a tensor to a Python number for ONNX convertation

I’m trying to convert lstm-decoder with beam search to onnx, but I got some TracerWarnings. Is it possible to operate with int-number item from 1-size input tensor in my future code like this?

def forward(self, img_features, beam_size):
    '''
    Arguments:
        img_features (tensor): Extracted features from the encoder module. (batch_size, feature_pixels = 7x7, encoder_dim = 2048)
        beam_size (tensor): Number of top candidates to consider for beam search. (int, = 3 or =6 or =9)
    '''
    item = beam_size.item()
    #some code...
    output = torch.tensor([item])
    return output

Now I have this warnings and constant output of Onnx-model:

TracerWarning: Converting a tensor to a Python number might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  item = beam_size.item()
TracerWarning: torch.tensor results are registered as constants in the trace. You can safely ignore this warning if you use this function to create tensors out of constant variables that would be the same every time you call this function. In any other case, this might cause the trace to be incorrect.
  output = torch.tensor([item])

I tried using @torch.jit.script but it didn’t help too

Could you post your real use case, i.e. how the Python integer is used?
Based on your current code snippet you should be able to just use the tensor instead of converting it to an integer, but I assume it’s just an example and not your real use case.

Ok, it’s just a bit of my problems with Onnx-convertation of lstm beam-search, and converting a tensor to a Python number is the first of them. When I export this for Onnx, I got the constant output from the model with incorrect tracing. I found, that I should script it with torch.script, but those examples not easy to understand :thinking:

    def forward(self, img_features, beam_size):
        """
        Function to generate the caption for the corresponding encoded image using beam search to provide the most optimal caption 
        combination. 
        
        Arguments:
            img_features (tensor): Extracted features from the encoder module. (batch_size, feature_pixels = 7x7, encoder_dim = 2048)
            beam_size (tensor): Number of top candidates to consider for beam search. (int = 3 or 6 or 9)
            
        Output:
            sentence (tensor): ordered list of words of the final optimal caption (list)

        """
        beam_size = beam_size.item()#!Converting a tensor to a Python number
        
        prev_words = torch.zeros(beam_size, 1).long()
        
        sentences = prev_words
        
        top_preds = torch.zeros(beam_size, 1)
        
        completed_sentences = []
        completed_sentences_preds = []
        
        step = 1
        h, c = self.get_init_lstm_state(img_features)

        while True:
            embedding = self.embedding(prev_words).squeeze(1)
            context = self.attention(img_features, h)[0]
            gate = self.sigmoid(self.f_beta(h))
            gated_context = gate * context

            lstm_input = torch.cat((embedding, gated_context), dim=1)
            h, c = self.lstm(lstm_input, (h, c))
            output = self.deep_output(h)
            output = top_preds.expand_as(output) + output

            if step == 1:
                top_preds, top_words = output[0].topk(beam_size, 0, True, True)
            else:
                top_preds, top_words = output.view(-1).topk(beam_size, 0, True, True)
            prev_word_idxs = top_words / output.size(1)
            next_word_idxs = top_words % output.size(1)

            sentences = torch.cat((sentences[prev_word_idxs], next_word_idxs.unsqueeze(1)), dim=1)

            incomplete = [idx for idx, next_word in enumerate(next_word_idxs) if next_word != 1]#!Converting a tensor to a Python boolean
            complete = list(set(range(len(next_word_idxs))) - set(incomplete))#!Converting a tensor to a Python index

            if len(complete) > 0:
                completed_sentences.extend(sentences[complete].tolist())#!Converting a tensor to a Python list
                completed_sentences_preds.extend(top_preds[complete])#!Converting a tensor to a Python index
            beam_size -= len(complete)

            if beam_size == 0:
                break
            sentences = sentences[incomplete]
            h = h[prev_word_idxs[incomplete]]
            c = c[prev_word_idxs[incomplete]]
            img_features = img_features[prev_word_idxs[incomplete]]
            top_preds = top_preds[incomplete].unsqueeze(1)
            prev_words = next_word_idxs[incomplete].unsqueeze(1)

            if step > 50:
                break
            step += 1

        idx = completed_sentences_preds.index(max(completed_sentences_preds))#!Converting a tensor to a Python boolean
        sentence = completed_sentences[idx]

        sentence = torch.tensor(sentence)#!torch.tensor results are registered as constants
        return sentence

Is scripting working for you now or are you still facing some issues?
I’m not an expert in using ONNX, but could you try to use tensors instead of Python literals via e.g.:

beam_size = torch.tensor([5])
prev_words = torch.zeros(beam_size, 1).long()

Yes, your case is working, but in this line, for example:

top_preds, top_words = output[0].topk(beam_size, 0, True, True)

‘topk()’ first argument must be int, not Tensor, so I still need this option (converatation tensor to a single number).
I try to script it like this:

        from torch import Tensor
        @torch.jit.script
        def get_item(x: Tensor):
            item = x.item()
            return item
        item = get_item(beam_size)

but got this error:

RuntimeError: get_item() Expected a value of type 'Tensor' for argument 'x' but instead found type 'int'.
Position: 0
Value: 1
Declaration: get_item(Tensor x) -> (Scalar)
Cast error details: Unable to cast Python instance of type <class 'int'> to C++ type 'at::Tensor'

What do you think about possible solution for this problems?

The last error message points towards an integer input, while a tensor is expected.
Did you accidentally pass x.item() to get_item()?

Sorry, I’m forgot to save changes in module :smile_cat:

Now, when I try to export to Onnnx this model with above get_item(beam_size) scripted function, I got
RuntimeError: Tracer cannot set value trace for type Int. Supported types are tensor, tensor list, and tuple of tensors.

Is it problem from Onnx-part, or am I doing something wrong with torch.script and it leads to errors?

I’m not sure, but the standalone PyTorch script seems to work:

from torch import Tensor
@torch.jit.script
def get_item(x: Tensor):
    item = x.item()
    return item

beam_size = torch.randn(1)
item = get_item(beam_size)
print(get_item.graph)
> graph(%x.1 : Tensor):
  %item.1 : Scalar = aten::item(%x.1) # <ipython-input-109-1655baf12d25>:4:11
  return (%item.1)
print(item)
> 0.18689797818660736
1 Like

Yes, this is work in practice, but not in onnx-export(

Anyway, thanks for your help, you are best :kissing_heart:

1 Like