diff --git a/.gitignore b/.gitignore index f7b6871..6056764 100644 --- a/.gitignore +++ b/.gitignore @@ -15,3 +15,4 @@ dataset/winter/ dataset/event.csv/ models/VPRTempo78415685001.pth models/VPRTempoQuant78415685001.pth +VPRTempo.egg-info/ diff --git a/README.md b/README.md index fc6a828..4244cda 100644 --- a/README.md +++ b/README.md @@ -20,7 +20,7 @@ In this repository, we provide two networks: To use VPRTempo, please follow the instructions below for installation and usage. -## :star: Update v1.1.1: What's new? +## :star: Update v1.1: What's new? - Full integration of VPRTempo into torch.nn architecture - Quantization Aware Training (QAT) enabled to train weights in int8 space - Addition of tutorials in Jupyter Notebooks to learn how to use VPRTempo as well as explain the computational logic diff --git a/dataset/test/images-00202.png b/dataset/test/images-00202.png new file mode 100644 index 0000000..56ae0ca Binary files /dev/null and b/dataset/test/images-00202.png differ diff --git a/main.py b/main.py index a8f745a..a5cb8c0 100644 --- a/main.py +++ b/main.py @@ -26,9 +26,7 @@ import argparse import sys sys.path.append('./src') -sys.path.append('./networks/base') -sys.path.append('./networks/quantized') - +sys.path.append('./vprtempo') import torch.quantization as quantization from VPRTempoTrain import VPRTempoTrain, generate_model_name, check_pretrained_model, train_new_model diff --git a/networks/base/VPRTempoTest.py b/networks/base/VPRTempoTest.py deleted file mode 100644 index cce50e0..0000000 --- a/networks/base/VPRTempoTest.py +++ /dev/null @@ -1,226 +0,0 @@ -#MIT License - -#Copyright (c) 2023 Adam Hines, Peter G Stratton, Michael Milford, Tobias Fischer - -#Permission is hereby granted, free of charge, to any person obtaining a copy -#of this software and associated documentation files (the "Software"), to deal -#in the Software without restriction, including without limitation the rights -#to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -#copies of the Software, and to permit persons to whom the Software is -#furnished to do so, subject to the following conditions: - -#The above copyright notice and this permission notice shall be included in all -#copies or substantial portions of the Software. - -#THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -#IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -#FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -#AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -#LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -#OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -#SOFTWARE. - -''' -Imports -''' - -import os -import torch -import gc -import sys -sys.path.append('./src') -sys.path.append('./models') -sys.path.append('./output') -sys.path.append('./dataset') - -import blitnet as bn -import numpy as np -import torch.nn as nn - -from dataset import CustomImageDataset, ProcessImage -from torch.utils.data import DataLoader -from tqdm import tqdm -from prettytable import PrettyTable -from metrics import recallAtK - -class VPRTempo(nn.Module): - def __init__(self, dims, args=None): - super(VPRTempo, self).__init__() - - # Set the arguments - if args is not None: - self.args = args - for arg in vars(args): - setattr(self, arg, getattr(args, arg)) - setattr(self, 'dims', dims) - if torch.cuda.is_available(): - self.device = "cuda:0" - else: - self.device = "cpu" - - # Set the dataset file - self.dataset_file = os.path.join('./dataset', self.dataset + '.csv') - - # Layer dict to keep track of layer names and their order - self.layer_dict = {} - self.layer_counter = 0 - - # Define layer architecture - self.input = int(self.dims[0]*self.dims[1]) - self.feature = int(self.input * 2) - self.output = int(args.num_places / args.num_modules) - - """ - Define trainable layers here - """ - self.add_layer( - 'feature_layer', - dims=[self.input, self.feature], - device=self.device, - inference=True - ) - self.add_layer( - 'output_layer', - dims=[self.feature, self.output], - device=self.device, - inference=True - ) - - def add_layer(self, name, **kwargs): - """ - Dynamically add a layer with given name and keyword arguments. - - :param name: Name of the layer to be added - :type name: str - :param kwargs: Hyperparameters for the layer - """ - # Check for layer name duplicates - if name in self.layer_dict: - raise ValueError(f"Layer with name {name} already exists.") - - # Add a new SNNLayer with provided kwargs - setattr(self, name, bn.SNNLayer(**kwargs)) - - # Add layer name and index to the layer_dict - self.layer_dict[name] = self.layer_counter - self.layer_counter += 1 - - def evaluate(self, models, test_loader, layers=None): - """ - Run the inferencing model and calculate the accuracy. - - :param test_loader: Testing data loader - :param layers: Layers to pass data through - """ - # Initialize the tqdm progress bar - pbar = tqdm(total=self.num_places, - desc="Running the test network", - position=0) - self.inferences = [] - for model in models: - self.inferences.append(nn.Sequential( - model.feature_layer.w, - nn.Hardtanh(0, 0.9), - nn.ReLU(), - model.output_layer.w, - nn.Hardtanh(0, 0.9), - nn.ReLU() - )) - # Initiliaze the output spikes variable - out = [] - # Run inference for the specified number of timesteps - for spikes, labels in test_loader: - # Set device - spikes, labels = spikes.to(self.device), labels.to(self.device) - # Forward pass - spikes = self.forward(spikes) - # Add output spikes to list - out.append(spikes.detach().cpu().tolist()) - pbar.update(1) - - # Close the tqdm progress bar - pbar.close() - - # Rehsape output spikes into a similarity matrix - out = np.reshape(np.array(out),(model.num_places,model.num_places)) - - # Recall@N - N = [1,5,10,15,20,25] # N values to calculate - R = [] # Recall@N values - # Create GT matrix - GT = np.zeros((model.num_places,model.num_places), dtype=int) - for n in range(len(GT)): - GT[n,n] = 1 - # Calculate Recall@N - for n in N: - R.append(round(recallAtK(out,GThard=GT,K=n),2)) - # Print the results - table = PrettyTable() - table.field_names = ["N", "1", "5", "10", "15", "20", "25"] - table.add_row(["Recall", R[0], R[1], R[2], R[3], R[4], R[5]]) - print(table) - - def forward(self, spikes): - """ - Compute the forward pass of the model. - - Parameters: - - spikes (Tensor): Input spikes. - - Returns: - - Tensor: Output after processing. - """ - - in_spikes = spikes.detach().clone() - outputs = [] # List to collect output tensors - - for inference in self.inferences: - out_spikes = inference(in_spikes) - outputs.append(out_spikes) # Append the output tensor to the list - - # Concatenate along the desired dimension - concatenated_output = torch.cat(outputs, dim=1) - - return concatenated_output - - def load_model(self, models, model_path): - """ - Load pre-trained model and set the state dictionary keys. - """ - combined_state_dict = torch.load(model_path, map_location=self.device) - - for i, model in enumerate(models): # models_classes is a list of model classes - model.load_state_dict(combined_state_dict[f'model_{i}']) - model.eval() # Set the model to inference mode - -def run_inference(models): - """ - Run inference on a pre-trained model. - - :param model: Model to run inference on - :param model_name: Name of the model to load - :param qconfig: Quantization configuration - """ - # Initialize the image transforms and datasets - image_transform = ProcessImage(models[0].dims, models[0].patches) - max_samples=models[0].num_places - - test_dataset = CustomImageDataset(annotations_file=models[0].dataset_file, - base_dir=models[0].data_dir, - img_dirs=models[0].query_dir, - transform=image_transform, - skip=models[0].filter, - max_samples=max_samples) - # Initialize the data loader - test_loader = DataLoader(test_dataset, - batch_size=1, - shuffle=False, - num_workers=8, - persistent_workers=True) - - # Retrieve layer names for inference - layer_names = list(models[0].layer_dict.keys()) - - # Use evaluate method for inference accuracy - with torch.no_grad(): - models[0].evaluate(models, test_loader, layers=layer_names) \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 60165d6..65baa9a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,5 @@ torch torchvision -torchaudio numpy pandas tqdm diff --git a/setup.py b/setup.py index c4a06b6..616f7b9 100644 --- a/setup.py +++ b/setup.py @@ -11,7 +11,6 @@ requirements = [ 'torch', 'torchvision', - 'torchaudio', 'numpy', 'pandas', 'tqdm', @@ -22,7 +21,7 @@ # define the setup setup( name="VPRTempo", - version="1.1.1", + version="1.1.2", description='VPRTempo: A Fast Temporally Encoded Spiking Neural Network for Visual Place Recognition', long_description=long_description, long_description_content_type='text/markdown', @@ -36,7 +35,7 @@ # 3 - Alpha # 4 - Beta # 5 - Production/Stable - 'Development Status :: 5 - Production/Stable', + 'Development Status :: 4 - Beta', # Indicate who your project is intended for 'Intended Audience :: Developers', diff --git a/test_vprtempo.py b/test_vprtempo.py deleted file mode 100644 index eb93dc3..0000000 --- a/test_vprtempo.py +++ /dev/null @@ -1,90 +0,0 @@ -#MIT License - -#Copyright (c) 2023 Adam Hines, Peter G Stratton, Michael Milford, Tobias Fischer - -#Permission is hereby granted, free of charge, to any person obtaining a copy -#of this software and associated documentation files (the "Software"), to deal -#in the Software without restriction, including without limitation the rights -#to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -#copies of the Software, and to permit persons to whom the Software is -#furnished to do so, subject to the following conditions: - -#The above copyright notice and this permission notice shall be included in all -#copies or substantial portions of the Software. - -#THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -#IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -#FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -#AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -#LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -#OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -#SOFTWARE. - -''' -Imports -''' -import argparse -import sys -sys.path.append('./src') -sys.path.append('./networks/base') -sys.path.append('./networks/quantized') - -import torch.quantization as quantization - -from VPRTempoTest import VPRTempo, run_inference - -def initialize_and_run_model(args,dims): - - models = [] - for _ in range(args.num_modules): - # Initialize the model - model = VPRTempo(dims, args) - models.append(model) - # Generate the model name - # Run the inference model - run_inference(models) - -def parse_network(): - ''' - Define the base parameter parser (configurable by the user) - ''' - parser = argparse.ArgumentParser(description="Args for base configuration file") - - # Define the dataset arguments - parser.add_argument('--dataset', type=str, default='nordland', - help="Dataset to use for training and/or inferencing") - parser.add_argument('--data_dir', type=str, default='./dataset/', - help="Directory where dataset files are stored") - parser.add_argument('--num_places', type=int, default=1, - help="Number of places to use for training and/or inferencing") - parser.add_argument('--num_modules', type=int, default=1, - help="Number of expert modules to use split images into") - parser.add_argument('--max_module', type=int, default=1, - help="Maximum number of images per module") - parser.add_argument('--database_dirs', nargs='+', default=['spring', 'fall'], - help="Directories to use for training") - parser.add_argument('--query_dir', nargs='+', default=['test'], - help="Directories to use for testing") - - # Define training parameters - parser.add_argument('--filter', type=int, default=8, - help="Images to skip for training and/or inferencing") - parser.add_argument('--epoch', type=int, default=4, - help="Number of epochs to train the model") - - # Define image transformation parameters - parser.add_argument('--patches', type=int, default=15, - help="Number of patches to generate for patch normalization image into") - parser.add_argument('--dims', type=str, default="56,56", - help="Dimensions to resize the image to") - - # Output base configuration - args = parser.parse_args() - dims = [int(x) for x in args.dims.split(",")] - - # Run the network with the desired settings - initialize_and_run_model(args,dims) - -def test_answer(): - # User input to determine if using quantized network or to train new model - parse_network() \ No newline at end of file diff --git a/networks/base/VPRTempo.py b/vprtempo/VPRTempo.py similarity index 100% rename from networks/base/VPRTempo.py rename to vprtempo/VPRTempo.py diff --git a/networks/quantized/VPRTempoQuant.py b/vprtempo/VPRTempoQuant.py similarity index 100% rename from networks/quantized/VPRTempoQuant.py rename to vprtempo/VPRTempoQuant.py diff --git a/networks/quantized/VPRTempoQuantTrain.py b/vprtempo/VPRTempoQuantTrain.py similarity index 100% rename from networks/quantized/VPRTempoQuantTrain.py rename to vprtempo/VPRTempoQuantTrain.py diff --git a/networks/base/VPRTempoTrain.py b/vprtempo/VPRTempoTrain.py similarity index 100% rename from networks/base/VPRTempoTrain.py rename to vprtempo/VPRTempoTrain.py diff --git a/vprtempo/__init__.py b/vprtempo/__init__.py new file mode 100644 index 0000000..8aafc8b --- /dev/null +++ b/vprtempo/__init__.py @@ -0,0 +1,2 @@ +from .VPRTempo import VPRTempo +from .VPRTempoQuantTrain import VPRTempoQuantTrain \ No newline at end of file