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!