Wrap model with extra hidden layer

Ben_Nicholl: <p>I have a pre trained transformer model that gets called like so</p>
<p><code>outputs = model(input_ids, attention_mask=mask, token_type_ids=tokens)</code></p>
<p>I then take the outputs of that model</p>
<p><code>last_hidden_state = outputs.last_hidden_state</code></p>
<p>and perform some mathematical manipulations, such as the ones below</p>
<p><code>last_hidden_state_zero_layer = torch.mul(last_hidden_state, some_parameter)</code><br>
<code>logits = logit_layer(last_hidden_state_zero_layer)</code></p>
<p>My question is, is there a way I can take the model and wrap the model plus the extra functions to simply create one PyTorch model? pseudocode below</p>
<pre><code class="lang-auto">`model = pytorch_wrapper_to_create_one_model = (

outputs = model(input_ids, attention_mask=mask, token_type_ids=tokens),

last_hidden_state_zero_layer = torch.mul(outputs.last_hidden_state, some_parameter),

logits = logit_layer(last_hidden_state_zero_layer)
)`
</code></pre>

thecho7: <p>Here’s a code snippet,</p>
<pre><code class="lang-auto">class yourModel(nn.Module):
    def __init__(self, ...):
        self.backbone = previous_model(...)
        self.last_hidden_state_zero_layer = nn.Linear(...)
        self.logit_layer = nn.Linear(...)

    def forward(self, input):
        outputs = self.backbone(input, ...)
        last_hidden = self.last_hidden_state_zero_layer(outputs.last_hidden_state, some_parameter)
        logit = self.logit_layer(last_hidden)
        return logit
</code></pre>
<p>Define a class with components is a possible approach</p>

Ben_Nicholl: <p>Thanks, works like a charm.</p>

Ben_Nicholl: <p>So this did work, but now when I try and deploy this on sagemaker, it is unable to load the model. When I look at the type of the above model, it reads<br>
<code>type(yourModel) = __main__.Model</code><br>
so I’m assuming this is because sagemaker doesn’t recognize this as a torch object. Do you know of a way to fix this?</p>

thecho7: <p>Could you share the code? Sagemaker is not a special something but just a instance (computer)</p>

Ben_Nicholl: <p>This is how I instantiate the model</p>
<pre><code class="lang-auto">val_to_col_elec_model = PyTorchModel(#model_data='s3://sagemaker-nl-filter-suggest-two/fields/pytorchmodel.tar.gz',
                         model_data='s3://natty-language-query/value_to_cols_model/model_.tar.gz',
                         role=role,
                         entry_point='torchserve.py',
                         source_dir='value_to_column_source_dir',
                         framework_version='1.6.0',
                         py_version='py36')
</code></pre>
<p>and this is how I deploy</p>
<pre><code class="lang-auto">nl_detector = val_to_col_elec_model.deploy(
                     serverless_inference_config=serverless_config,
                     initial_instance_count = 1,
                     serializer=JSONSerializer(),
                     deserializer=JSONDeserializer(),
                     instance_type = 'ml.t2.medium', endpoint_name = endpoint_name)
</code></pre>
<p>When I upload the model created with this code</p>
<pre><code class="lang-auto">class Model(torch.nn.Module):
    def __init__(self, model, logit_layer):
        super(Model, self).__init__()
        
        self.model = model
        self.model.eval()
        #self.last_hidden_state_zero_layer = nn.Linear(...)
        self.logit_layer = logit_layer
        self.logit_layer.eval()

    def forward(self, input_ids, mask, tokens):
        pre_logits_mask = torch.reshape(mask, (mask.shape[0], mask.shape[1], 1) )
        
        outputs = self.model(input_ids, attention_mask=mask, token_type_ids=tokens)
        last_hidden_state = outputs.last_hidden_state
        last_hidden_state_zero_layer = torch.mul(last_hidden_state, pre_logits_mask)
        summed_final_hidden_state = torch.sum(last_hidden_state_zero_layer, 1)
        
        logits = self.logit_layer(summed_final_hidden_state)
        probs = torch.sigmoid(logits)
        
        return probs
    
entire_model = Model(model, logit_layer)
</code></pre>
<p>It runs into a model not loading error. again I believe this is because it is of type <strong>main</strong>.Model, because the sagemaker container is expecting a torch object.</p>
<p>I tested this by uploading only the sub model that I wrapped in the above Model class, which is of type transformers.models.electra.modeling_electra.ElectraModel and it worked.</p>
<p>Also, the above models save and load perfectly fine when ran locally. Let me know if anything needs any clarification</p>

thecho7: <p>Doesn’t Sagemaker need 4 functions for deployment?<br>
I wrote <code>input_fn</code>, <code>output_fn</code>, <code>model_fn</code>, <code>predict_fn</code> when I made a service.<br>
I don’t get how did ElectraModel run successfully but would you change the name of class from <code>Model</code> to something different? I’m pretty not sure it helps but this is my tiny guess that the name of <code>Model</code> is conflicted to somewhat from Sagemaker class.</p>
<p>Out of Sagemaker, you don’t need to call <code>.eval()</code> separately.<br>
Using <code>entire_model.eval()</code> is enough.</p>

Ben_Nicholl: <p>yes, I have all of the logic pertaining to the <code>input_fn</code> , <code>output_fn</code> , <code>model_fn</code> , <code>predict_fn</code> functions in the source directory. <code>source_dir='value_to_column_source_dir'</code> points to that logic.</p>
<p>I changed the name of the class to something other than Model and it still doesn’t work. I wonder if anyone else had tried deploying a PyTorch model created this way using sagemaker.</p>

thecho7: <p>Please let me know error messages from logs</p>

Ben_Nicholl: <p>So I’m using <code>torch.load</code> in my <code>model_fn</code> function, and this is the error I’m getting in CloudWatch. Also, I renamed <code>class Model(torch.nn.Module)</code> to <code>class transformer_Model(torch.nn.Module)</code></p>
<blockquote>
<p>2022-07-18 01:58:04,499 [INFO ] W-9000-model_1-stdout org.pytorch.serve.wlm.WorkerLifeCycle - AttributeError: Can’t get attribute ‘transformer_Model’ on &lt;module ‘<strong>main</strong>’ from ‘/opt/conda/lib/python3.6/site-packages/ts/model_service_worker.py’&gt;</p>
</blockquote>

Ben_Nicholl: <p>It turns out you have to save the state_dict instead of the entire model. Then when you want to use that custom model, you import the <code>class Model(torch.nn.Module)</code>, then call <code>Model.load_state_dict(torch.load('path'))</code> to fill in the fine tuned weights.</p>
<p><a href="https://pytorch.org/tutorials/beginner/saving_loading_models.html#save-load-entire-model" class="onebox" target="_blank" rel="noopener nofollow ugc">https://pytorch.org/tutorials/beginner/saving_loading_models.html#save-load-entire-model</a></p>

thecho7: <p>Nicely done:)<br>
It was so hard to recognize</p>