diff --git a/tools/gen_random_vector.py b/tools/gen_random_vector.py index 6a6fc2b..cf96820 100644 --- a/tools/gen_random_vector.py +++ b/tools/gen_random_vector.py @@ -2,9 +2,11 @@ import numpy as np import utils from sys import argv + def gen_random_vector(size): return np.random.rand(size) + if __name__ == "__main__": if len(argv) < 2: @@ -12,4 +14,4 @@ if __name__ == "__main__": exit(1) vector = gen_random_vector(int(argv[1])) - utils.print_cpp_vector(vector) \ No newline at end of file + utils.print_cpp_vector(vector) diff --git a/tools/utils.py b/tools/utils.py index 95da2b8..25a93fd 100644 --- a/tools/utils.py +++ b/tools/utils.py @@ -1,7 +1,9 @@ import torch import struct -import numpy as np +from PIL import Image + +from torchvision import transforms def print_cpp_vector(vector, name="expected"): @@ -14,7 +16,7 @@ def print_cpp_vector(vector, name="expected"): def export_model_weights(model: torch.nn.Module, filename): - with open(filename, 'wb') as f: + with open(filename, "wb") as f: version = 1 header = "" @@ -22,23 +24,54 @@ def export_model_weights(model: torch.nn.Module, filename): tensor_data = b"" for name, param in model.named_parameters(): - if 'weight' not in name and 'bias' not in name: + if "weight" not in name and "bias" not in name: continue tensor_bytes = param.type(torch.float32).detach().numpy().tobytes() tensor_size = param.numel() - header += f"{name},{tensor_size},{offset}\n" + header += f"{name},{tensor_size},{offset}\n" offset += len(tensor_bytes) tensor_data += tensor_bytes f.seek(0) - f.write(struct.pack('H', version)) - f.write(struct.pack('Q', len(header))) - f.write(header.encode('utf-8')) + f.write(struct.pack("H", version)) + f.write(struct.pack("Q", len(header))) + f.write(header.encode("utf-8")) f.write(tensor_data) + def print_model_parameters(model: torch.nn.Module): for name, param in model.named_parameters(): print(name, param.numel()) + + +def predict(model, image_path, preprocess=None): + input_image = Image.open(image_path) + + if preprocess is None: + preprocess = transforms.Compose( + [ + transforms.Resize(299), + transforms.CenterCrop(299), + transforms.ToTensor(), + transforms.Normalize( + mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225] + ), + ] + ) + + input_tensor = preprocess(input_image) + input_batch = input_tensor.unsqueeze( + 0 + ) # create a mini-batch as expected by the model + + # move the input and model to GPU for speed if available + if torch.cuda.is_available(): + input_batch = input_batch.to("cuda") + model.to("cuda") + + with torch.no_grad(): + output = model(input_batch) + return torch.argmax(output)