diff --git a/benchmark/torchscripts/BernoulliCRS.pt b/benchmark/torchscripts/BernoulliCRS.pt new file mode 100644 index 00000000..e635a5f3 Binary files /dev/null and b/benchmark/torchscripts/BernoulliCRS.pt differ diff --git a/benchmark/torchscripts/BernoulliCRS.py b/benchmark/torchscripts/BernoulliCRS.py new file mode 100644 index 00000000..25b5859a --- /dev/null +++ b/benchmark/torchscripts/BernoulliCRS.py @@ -0,0 +1,66 @@ +import torch +import time +import os + +def get_first_element(tensor): + if tensor.numel() == 1: + return tensor.item() + else: + return tensor[0].item() + +def is_empty_tensor(tensor): + return tensor.numel() == 0 + +@torch.jit.script +def BernoulliCRS(A: torch.Tensor, B: torch.Tensor, k: int): + # Get the dimension of A + A = A.t() + n, m = A.shape + + assert n == B.shape[0] + assert k < n + + # probability distribution + sample = torch.rand(n) # default: uniform + sample = torch.div(sample, sample.sum() / k) # sum = k as per the paper + + # diagonal scaling matrix P (nxn) + P = torch.diag(1.0 / torch.sqrt(sample)) + + # random diagonal sampling matrix K (nxn) + sample = (torch.rand(n) < sample).float() + K = torch.diag(sample) + + a = torch.matmul(torch.matmul(A.t(), P), K) + b = torch.matmul(torch.matmul(a, K), P) + + return torch.matmul(b, B) + + +def main(): + + width = 1000 + A = torch.rand(2000, width) + B = torch.rand(width, 2000) + + t = time.time() + + aResult = BernoulliCRS(A, B, 500) + print("approximate: " + str(time.time() - t) + "s") + + print(aResult) + + # exact result + t = time.time() + eResult = torch.matmul(A, B) + print("\nExact: " + str(time.time() - t) + "s") + + print(eResult) + + print("\nerror: " + str(torch.norm(aResult - eResult, p='fro').item())) + + script = BernoulliCRS.save("BernoulliCRS.pt") + +if __name__ == '__main__': + main() + \ No newline at end of file diff --git a/benchmark/torchscripts/CRS.pt b/benchmark/torchscripts/CRS.pt new file mode 100644 index 00000000..b405bcb9 Binary files /dev/null and b/benchmark/torchscripts/CRS.pt differ diff --git a/benchmark/torchscripts/CRSV2.pt b/benchmark/torchscripts/CRSV2.pt new file mode 100644 index 00000000..12de4027 Binary files /dev/null and b/benchmark/torchscripts/CRSV2.pt differ diff --git a/benchmark/torchscripts/ColumnRowSampling.py b/benchmark/torchscripts/ColumnRowSampling.py new file mode 100644 index 00000000..9e7d57c3 --- /dev/null +++ b/benchmark/torchscripts/ColumnRowSampling.py @@ -0,0 +1,70 @@ +import torch +import time +import os +import math + +def get_first_element(tensor): + if tensor.numel() == 1: + return tensor.item() + else: + return tensor[0].item() + +def is_empty_tensor(tensor): + return tensor.numel() == 0 + +@torch.jit.script +def CRS(A: torch.Tensor, B: torch.Tensor, k: int): + # Get the dimension of A + A = A.t() + n, m = A.shape + + assert n == B.shape[0] + assert k < n + + # probability distribution + probs = torch.ones(n) / n # default: uniform + + # sample k indices from range 0 to n for given probability distribution + indices = torch.multinomial(probs, k, replacement=True) + + # Sample k columns from A + A_sampled = A[indices, :] + ratio = math.ceil(n / k) + A_sampled = torch.div((A_sampled / k).t(), probs[::ratio]) + + # Sample k rows from B + B_sampled = B[indices, :] + + # Compute the matrix product + result = A_sampled.matmul(B_sampled) + return result + + +def main(): + + width = 2000 + A = torch.rand(5000, width) + B = torch.rand(width, 5000) + + + t = time.time() + + aResult = CRS(A, B, 500) + print("approximate: " + str(time.time() - t) + "s") + + print(aResult) + + # exact result + t = time.time() + eResult = torch.matmul(A, B) + print("\nExact: " + str(time.time() - t) + "s") + + print(eResult) + + print("\nerror: " + str(torch.norm(aResult - eResult, p='fro').item())) + + FDAMM_script = CRS.save("CRS.pt") + +if __name__ == '__main__': + main() + \ No newline at end of file diff --git a/benchmark/torchscripts/ColumnRowSamplingVer2.py b/benchmark/torchscripts/ColumnRowSamplingVer2.py new file mode 100644 index 00000000..ef67b4e2 --- /dev/null +++ b/benchmark/torchscripts/ColumnRowSamplingVer2.py @@ -0,0 +1,71 @@ +import torch +import time +import os + +def get_first_element(tensor): + if tensor.numel() == 1: + return tensor.item() + else: + return tensor[0].item() + +def is_empty_tensor(tensor): + return tensor.numel() == 0 + +@torch.jit.script +def CRS(A: torch.Tensor, B: torch.Tensor, k: int): + # Get the dimension of A + A = A.t() + n, m = A.shape + + assert n == B.shape[0] + assert k < n + + # probability distribution + # dist = torch.distributions.Uniform(0, 1) + sample = torch.rand(n) # default: uniform + + # diagonal scaling matrix D (nxn) + sample = torch.div(sample, sample.sum()) + D = torch.diag(1.0 / torch.sqrt(k * sample)) + + # sampling matrix S (kxn) + column_indices = torch.multinomial(sample, k, replacement=True) + S = torch.zeros(k, n) + for row, col in enumerate(column_indices): + S[row, col] = 1 + + a = torch.matmul(torch.matmul(A.t(), D), S.t()) + b = torch.matmul(torch.matmul(a, S), D) + + + return torch.matmul(b, B) + + +def main(): + + + width = 2000 + A = torch.rand(1000, width) + B = torch.rand(width, 1000) + + t = time.time() + + aResult = CRS(A, B, 500) + print("approximate: " + str(time.time() - t) + "s") + + print(aResult) + + # exact result + t = time.time() + eResult = torch.matmul(A, B) + print("\nExact: " + str(time.time() - t) + "s") + + print(eResult) + + print("\nerror: " + str(torch.norm(aResult - eResult, p='fro').item())) + + script = CRS.save("CRSV2.pt") + +if __name__ == '__main__': + main() + \ No newline at end of file