Untitled

model.load_state_dict(torch.load("mnist_cnn.pt"))
weight1 = model.conv1.weight.int().tolist()
weight2 = model.conv2.weight.int().tolist()
weight3 = model.fc1.weight.int().T.tolist()
weight4 = model.fc2.weight.int().T.tolist()

@fhe.compiler({"x": "encrypted"})
def f(x):
    x = fhe.conv(x, weight1, kernel_shape=(3, 3), strides=(1, 1))
    x = fhe.relu(x)
    x = fhe.conv(x, weight2, kernel_shape=(3, 3), strides=(1, 1))
    x = fhe.relu(x)
    x = fhe.maxpool(x, kernel_shape=(2, 2), strides=(2, 2))
    x = x.reshape(x.shape[0], -1)
    x = x.dot(weight3)
    x = fhe.relu(x)
    x = x.dot(weight4)
    return x

inputset = [np.random.randint(0, 255, size=(1, 1, 28, 28)) for _ in range(10)]
circuit = f.compile(inputset, relu_on_bits_threshold=9, use_gpu=True)

sample = np.random.randint(0, 255, size=(1, 1, 28, 28))
assert np.array_equal(circuit.encrypt_run_decrypt(sample), f(sample))