Tracing Python Objects in torch.jit.tracing for save and load

Viewed 329

I am trying to send a pytorch model from one server to another. I chose to use torch.jit.tracing to save the model for transmission. Now, the problem is that, I want to send encrypted parameters, which is a list of int python objects (to save 100+ digits for each ciphertext). I tried to save it in a dict() object and trace it.

# Model
class Net(nn.Module):
    def __init__(self, input_size, output_size):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(input_size, 50)
        # nn.init.normal_(self.fc1.weight, mean=0, std=1)
        self.fc2 = nn.Linear(50, 10)
        # nn.init.normal_(self.fc2.weight, mean=0, std=1)
        self.fc3 = nn.Linear(10, output_size)
        nn.init.normal_(self.fc3.weight, mean=0, std=1)
        self.fc1_enc = None
        self.fc2_enc = None
        self.fc3_enc = None
        # a dict of python int objects to save encrypted big ints.
        self.encrypted_params = {"int_wei": [[74812937489217492138741294723198479123749321874,  
                                              11111112222222222222888888888888888888888888888]]}
        self.exponent_encrypted_params = dict()

    def forward(self, x):
        x = F.elu(self.fc1(x))
        x = F.elu(self.fc2(x))
        return F.log_softmax(self.fc3(x), dim=-1)

And I trace the model by :

    model = Net(119, 2).to(device)
    traced_model = torch.jit.trace(model, torch.zeros(119, dtype=torch.float))

However, I failed to load the dict wit error: AttributeError: 'ScriptModule' object has no attribute 'encrypted_params'.

How can I serialize the model with big integer objects as parameters? Thanks!

0 Answers
Related