modelmark 0.1.0__py3-none-any.whl

modelmark/common/tester.py

@@ -0,0 +1,162 @@
+from __future__ import annotations
+
+import numpy as np
+import time
+import torch
+import torch.nn as nn
+from thop import profile
+from torch.utils.data import DataLoader
+from torch.nn.utils import clip_grad_norm_
+from modelmark.common.loader import Loader
+
+import logging
+logger = logging.getLogger(__name__)
+from modelmark.common.utils import load_config
+config = load_config()
+
+class Tester:
+	"""This class runs a benchmark of the current model and loader"""
+
+	def __init__(self, model : nn.Module, loader : Loader):
+
+		self.model = model
+		self.train_loader = loader.train_loader
+		self.val_loader = loader.val_loader
+		self.test_loader = loader.test_loader
+
+		self.optimizer = config.test_config["optim"]["class"](self.model.parameters(), lr=config.test_config["optim"]["lr"])
+		self.criterion = config.test_config["optim"]["criterion"]
+		self.metrics = [v for _, v in config.test_config["metrics"].items()]
+		self.num_epochs = config.test_config["optim"]["num_epochs"]
+		self.max_norm = config.test_config["optim"]["max_norm"]
+				
+		self.device = config.test_config["optim"]["device"]
+
+	def test(self):
+		"""
+		Train, validate, return to the best model weight state according to the validation loss, test and return.
+		"""
+
+		best_val_loss = float("inf")
+		best_state = None
+		train_time, train_gflops, train_mem = None, None, None
+
+		for epoch in range(self.num_epochs):
+
+			results = self.train(self.train_loader, measure_time=True, measure_flops=True, measure_memory=True)
+
+			train_loss = results[0]
+			train_time = results[1]
+			train_gflops = results[2]
+			train_mem = results[3]
+
+			val_loss = self.evaluate(self.val_loader, self.criterion)
+		
+			logger.debug(f"Epoch {epoch:03d} | T {train_loss:.4f} | V {val_loss:.4f}")
+		
+			# Save best model according to validation performance
+			if val_loss < best_val_loss:
+				best_val_loss = val_loss
+		
+				best_state = {
+					k: v.detach().cpu().clone()
+					for k, v in self.model.state_dict().items()
+				}
+		
+		# Restore best checkpoint
+		self.model.load_state_dict(best_state)
+
+		# Get final test metrics values
+		test_metrics = np.array([self.evaluate(self.test_loader, m) for m in self.metrics])
+		
+		return test_metrics, train_time, train_gflops, train_mem
+
+	def train(self, loader: DataLoader, measure_time: bool = False,
+			  measure_flops: bool = False, measure_memory: bool = False):
+		"""Train the model for one epoch."""
+
+		self.model.train()
+
+		total_loss = 0.0
+		total_samples = 0
+		total_time = 0
+		total_flops = 0
+		
+		is_cuda = self.device.type == "cuda"
+
+		if measure_memory and is_cuda:
+			torch.cuda.reset_peak_memory_stats(self.device)
+
+		if measure_time:
+			if is_cuda:
+				torch.cuda.synchronize()
+			start_time = time.perf_counter()
+
+		for i, (x, y) in enumerate(loader):
+
+			x = x.to(self.device)
+			y = y.to(self.device)
+
+			if measure_flops and (i == 0):
+				with torch.no_grad():
+					flops_per_batch, _ = profile(self.model, inputs=(x,), verbose=False)
+				flops_per_batch *= 3
+
+			self.optimizer.zero_grad(set_to_none=True)
+			pred = self.model(x)
+			loss = self.criterion(pred.squeeze(), y.squeeze())
+			loss.backward()
+			clip_grad_norm_(self.model.parameters(), max_norm=self.max_norm)
+			self.optimizer.step()
+
+			batch_size = x.size(0)
+			total_loss += loss.item() * batch_size
+			total_samples += batch_size
+
+			if measure_flops:
+				total_flops = flops_per_batch
+
+		avg_loss = total_loss / total_samples
+
+		results = [avg_loss]
+
+		if measure_time:
+			if is_cuda:
+				torch.cuda.synchronize()
+			total_time = time.perf_counter() - start_time
+			results.append(total_time)
+
+		if measure_flops:
+			results.append(total_flops / 1e9)
+
+		if measure_memory:
+			if is_cuda:
+				peak_mem_gb = torch.cuda.max_memory_allocated(self.device) / (1024 ** 3)
+			else:
+				peak_mem_gb = 0.0  # or float('nan') to signal "not applicable"
+			results.append(peak_mem_gb)
+
+		return tuple(results)
+
+	@torch.no_grad()
+	def evaluate(self, loader: DataLoader, criterion):
+		"""Evaluate the model on provided data"""
+	
+		self.model.eval()
+	
+		total_loss = 0.0
+		total_samples = 0
+	
+		for x, y in loader:
+			x = x.to(self.device)
+			y = y.to(self.device)
+	
+			pred = self.model(x)
+			loss = criterion(pred.squeeze(), y.squeeze())
+	
+			batch_size = x.size(0)
+			total_loss += loss.item() * batch_size
+			total_samples += batch_size
+	
+		return total_loss / total_samples
+	

modelmark/common/utils.py

@@ -0,0 +1,105 @@
+from __future__ import annotations
+
+
+import torch
+import torch.nn as nn
+
+import os
+import sys
+import time
+import random
+import numpy as np
+
+import importlib.util
+import modelmark.constants as constants
+
+import logging
+logger = logging.getLogger(__name__)
+
+def set_seed(s : int = 0):
+	# Seed all the generators for reproducible results
+	torch.random.manual_seed(s)
+	torch.cuda.manual_seed(s)
+	torch.cuda.manual_seed_all(s)
+	torch.use_deterministic_algorithms(True)
+	torch.backends.cudnn.deterministic = True
+	torch.backends.cudnn.benchmark = False
+	np.random.seed(s)
+	random.seed(s)
+
+def count_parameters(module: nn.Module) -> int:
+	"""Count the total number of elements/parameters in any PyTorch module."""
+	return sum(p.numel() for p in module.parameters()) / 1000
+
+# Setup non-blocking key input based on Operating System
+if os.name == 'nt':
+	import msvcrt
+	def get_keypress():
+		if msvcrt.kbhit():
+			# Return decoded string character
+			return msvcrt.getch().decode('utf-8', errors='ignore')
+		return None
+else:
+	import select
+	import termios
+	import tty
+	def get_keypress():
+		# Check if stdin has data waiting
+		if select.select([sys.stdin], [], [], 0)[0]:
+			fd = sys.stdin.fileno()
+			old_settings = termios.tcgetattr(fd)
+			try:
+				tty.setraw(sys.stdin.fileno())
+				ch = sys.stdin.read(1)
+			finally:
+				termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)
+			return ch
+		return None
+
+def start_timer(duration_seconds):
+	start_time = time.time()
+	#print("Timer started. Press any key to interrupt...")
+	
+	while True:
+		elapsed = time.time() - start_time
+		remaining = max(0, duration_seconds - elapsed)
+		
+		sys.stdout.write(f"\rStart in: {remaining:.1f}s")
+		sys.stdout.flush()
+		
+		# Check for user interruption
+		pressed_key = get_keypress()
+		if pressed_key is not None:
+			print("\n")
+			return pressed_key
+			
+		if remaining <= 0:
+			break
+			
+		time.sleep(0.05) # Lower sleep window for snappier key detection
+
+	print("\n")
+	#print("\n\nTimer finished naturally!")
+	return None
+
+def load_config():
+	"""Load the config (user or pkg)"""
+
+	user_config_path = constants.USER_CONFIG_PATH
+
+	if user_config_path.exists():
+	
+		# Add the user working dir to sys.path
+		user_dir = str(constants.USER_DIR)
+		if user_dir not in sys.path:
+			sys.path.insert(0, user_dir)
+
+		spec = importlib.util.spec_from_file_location("config", user_config_path)
+		user_config = importlib.util.module_from_spec(spec)
+		spec.loader.exec_module(user_config)
+		return user_config
+
+	else:
+	
+		from modelmark import config
+		return config

modelmark/config.py

@@ -0,0 +1,97 @@
+"""This is the testing config file example, customize as you need."""
+
+import torch
+import torch.nn as nn
+
+# ---- Model import secion ---- #
+# Put your model class definitions here
+# Either built-in:
+from modelmark.models.gru import GRUModel
+from modelmark.models.conv import ConvModel
+from modelmark.models.lstm import LSTMModel
+from modelmark.models.linear import Linear
+# Either custom (uncomment)
+#from models.linear import Linear
+
+# ---- Model configuration section ----
+model_config = {
+	"Linear" : {
+		"num_layers": 3,
+		"hidden_size": 64,
+		"class": Linear
+	},
+	"Conv" : {	
+		"num_layers": 3,
+		"hidden_size": 64,
+		"kernel_size": 4,
+		"class": ConvModel,
+	},
+	"GRU" : {	
+		"num_layers": 2,
+		"hidden_size": 64,
+		"class": GRUModel,
+	},
+	"LSTM" : {	
+		"num_layers": 2,
+		"hidden_size": 64,
+		"class": LSTMModel,
+	},
+}
+
+# ---- Data configuration section ----
+# Your datasets folder path
+data_path = "data/"
+# Configuration of your datasets
+data_config = {
+	"ETTh1" : {
+		"path": "ett/ETTh1.csv",  
+		"input_features": ["HUFL", "HULL", "MUFL", "MULL", "LUFL", "LULL", "OT"], 
+		"output_features": ["OT"],
+		"train_ratio": 0.6,
+		"val_ratio": 0.2, 
+	},
+	"ETTh2" : {
+		"path": "ett/ETTh2.csv",  
+		"input_features": ["HUFL", "HULL", "MUFL", "MULL", "LUFL", "LULL", "OT"], 
+		"output_features": ["OT"],
+		"train_ratio": 0.6,
+		"val_ratio": 0.2,
+	},
+}
+
+# ---- Testing configuration section ----
+
+# You can define your own metric
+import torch
+import torch.nn.functional as F
+def rmse(x, y):
+	return torch.sqrt(F.mse_loss(x, y))
+from torch.optim import Adam
+
+test_config = {
+	"optim": {
+		"name" : "Adam",
+		"class": Adam,
+		"max_norm": 1.0,
+		"criterion": nn.MSELoss(),
+		"lr": 1e-3,
+		"batch_size": 512,
+        "num_epochs": 5,
+		"device": torch.device("cuda"),
+	},
+	"metrics":
+	{
+		"RMSE": rmse,
+		"MAE": nn.L1Loss(),
+	},
+	"seeds": [1, 2, 3],
+	"contexts": [96, 192, 336, 720],
+	}
+
+# Other
+_data_split = ", ".join([f"{k} ({v["train_ratio"]:.2f}/{v["val_ratio"]:.2f}/{(1 - (v["train_ratio"] + v["val_ratio"])):.2f})" for k, v in data_config.items()])
+# --- Report configuration section ---
+title = f"Models evaluation report."
+subtitle=f"""Models were trained with {test_config["optim"]["name"]} optimizer, batch size = {test_config["optim"]["batch_size"]}, LR = {test_config["optim"]["lr"]}, epochs = {test_config["optim"]["num_epochs"]}
+Dataset split (train/val/test): {_data_split}
+Results are mean values over {len(test_config["seeds"])} runs, seeds used: {", ".join(map(str, test_config["seeds"]))}"""

modelmark/constants.py

@@ -0,0 +1,34 @@
+from pathlib import Path
+
+# ModelMark pkg dir path
+PACKAGE_DIR = Path(__file__).parent
+# User cwd path
+USER_DIR = Path.cwd()
+
+# Create user config dir (if not already exists)
+USER_CONFIG_DIR = USER_DIR / "modelmark_files"
+
+# Pkg config file path
+PACKAGE_CONFIG_PATH = PACKAGE_DIR / "config.py"
+# User config file path
+USER_CONFIG_PATH = USER_CONFIG_DIR / "config.py"
+
+# Create user models dir (if not already created)
+USER_MODELS_DIR = USER_DIR / "models"
+
+# Pkg example model path
+PACKAGE_MODEL_PATH = PACKAGE_DIR / "models" / "linear.py"
+# User example mopdel path
+USER_MODEL_PATH = USER_DIR / "models" / "linear.py"
+
+PARSER_DESC = """MODELMARK
+
+	Python tool to measure the performance of a custom neural network model
+	and compare it to other popular architectures."""
+PARSER_HELP = """\nModelmark available TASKs:
+init - Test init, create the modelmark_files/, customizable config.py file and models/ with examples.
+load - Dataset download, create the data/ folder and download ETT dataset there.
+run - Run the test.
+clear - Remove the modelmark_files folder.
+reset - Runs clean and init.
+"""

modelmark/modelmark.py

@@ -0,0 +1,206 @@
+"""
+ModelMark
+
+	This tool will help you to test NN models against each other,
+	and form a detailed report that is easy to embed to a website.   
+
+	It can download dataset with
+		modelmark -t load
+	It can run the test with 
+		modelmark -t run
+
+	Test consists of F * C * M * S steps, where:
+		F - number of dataset files in the config (e.g. ["ETTh1" : ..., "Weather" : ...] means F = 2)
+		C - number of context sizes (e.g. [32, 64, 128] means C = 3)
+		M - number of models (e.g. ["Linear" : ..., "LSTM" : ...] means M = 2)
+		S - number of seeds (e.g. [42, 43, 44] means S = 2)
+
+	At each testing run iteration, modelmark:
+		
+		1) Selects next dataset, context, model and seed 
+		2) Seeds the generators for reproducibility
+		3) Creates the loader, model and tester objects
+		4) Trains the model for E epochs, restores the state with the least validation loss
+		5) Tracks the GFLOPs (AVG over one batch), Memory (AVG peak usage during full training), Time (AVG per epoch)
+		5) Evaluates the model on metrics from configuration file (config.test_metric1&2)
+		6) Stores the mean result over S runs 
+
+		That way, the more seeds you run, the more "fair" the results are.
+	
+	Finally, modelmark will form the report with all the testing results, training stats and your machine metadata.
+
+	About training stats:
+		
+		Time - average time per epoch 
+		Params - total number of model params
+		GFLOPs - average per batch
+		Peak Memory - max per training iteration
+
+	---	
+	Example of the model: "modelmark/models/linear.py" 
+	Example of the config: "modelmark/config.py"
+"""
+
+import sys
+import torch
+import numpy as np
+import pandas as pd
+from tqdm import tqdm
+
+# Import custom functions
+from modelmark.common.utils import  set_seed, count_parameters, start_timer
+from modelmark.common.logger import setup_logging
+# Import custom classes
+from modelmark.common.report import Report
+from modelmark.common.parser import Parser
+from modelmark.common.loader import Loader
+from modelmark.common.tester import Tester
+
+import logging
+logger = logging.getLogger(__name__)
+from rich.console import Console
+console = Console()
+from modelmark.common.utils import load_config
+config = load_config()
+
+def _check_model(name = "Linear", device=config.test_config["optim"]["device"]):
+	"""Check the testing model on shapes compatibility."""
+	first_config = next(iter(config.data_config.values()))
+
+	input_dim = len(first_config["input_features"])
+	output_dim = len(first_config["output_features"])
+
+	model = config.model_config[name]["class"](input_dim, output_dim, context_size = 10).to(device)
+	dummy_input = torch.zeros([4, 10, input_dim]).to(device)
+	dummy_target = torch.zeros([4, 10, output_dim]).to(device)
+	generated = model(dummy_input)
+	if generated.shape == dummy_target.shape:
+		console.print(f"Testing model {name} is ready", style="green")
+		return None
+	else:
+		console.print(f"Testing model output expected shape {dummy_target.shape}, got {generated.shape}.\nCheck your definition of {name} model.", style="red")
+		return 1
+
+def run() -> int:
+
+	# Parse the arguments
+	parser = Parser()
+	code = parser.run()
+	if code is not None:
+		return code
+
+	# Setup logger
+	setup_logging()
+	logger.info("Logger started")
+		
+	# Check the models
+	for k, _ in config.model_config.items():
+		if _check_model(k) is not None:
+			return 1
+		
+	# Display test information 
+	logger.info("Display test info.")
+	console.print("-" * 128, style = "cyan")
+	console.print("Info\n", style = "magenta")
+	console.print(f"This test runs multiple training/validation/testing passes of all configured model architectures: {[k for k, _ in config.model_config.items()]}", style = "white")
+	console.print(f"To provide you with the most fair comparision, we run multiple tests over seeds, and average the results.", style = "white")
+	console.print(f"There is no parallel execution at the moment, so the test may take a while to finish (it depends on your config and hardware).", style = "white")
+	console.print("-" * 128, style="cyan")
+
+	console.print("Start the test? [y/n] (Will start automatically in 30s.)", style = "white")
+	k = start_timer(30)
+	if (k == 'n') or (k == 'N'):
+		logger.error("Test terminated by user.")
+		console.print("Terminated", style="red")
+		return 1
+
+	# Start the test
+	logger.info("Starting the test.")
+	console.print("Starting the test.", style = "green")
+
+	# Create test buffers
+	report_data = []
+	report_dataset_names = list(config.data_config)
+	report_model_names = list(config.model_config)
+	report_metric_names = list(config.test_config["metrics"])
+	report_stats = np.zeros((len(report_model_names), 3))
+	report_params = []
+
+	device = config.test_config["optim"]["device"]
+	c_current, c_total = 1, len(report_dataset_names) * len(report_model_names) * len(config.test_config["seeds"]) * len(config.test_config["contexts"])
+	pbar = tqdm(total=c_total, desc="Progress")
+	for file_name, file_config in config.data_config.items():
+
+		for test_context in config.test_config["contexts"]:
+
+			report_line = []
+			for model_name, model_config in config.model_config.items():
+
+				total_test_metrics = np.zeros(len(config.test_config["metrics"]))
+				for test_seed in config.test_config["seeds"]:
+					
+					# (Re-)Seed
+					set_seed(test_seed)
+					tqdm.write(f"Running test ({c_current}/{c_total}) | File: {file_name} | Context: {test_context} | Model: {model_name} | Seed: {test_seed}")
+					logger.debug(f"Test ({c_current}/{c_total})| F={file_name} C={test_context} M={model_name} S={test_seed}")
+				
+					# Create loader
+					loader = Loader(file_config = file_config, context_size = test_context)
+					# Create model
+					input_dim = len(file_config["input_features"])
+					output_dim = len(file_config["output_features"])
+					model = model_config["class"](input_dim, output_dim, context_size = test_context).to(device)
+					# Create tester
+					tester = Tester(model = model, loader = loader)
+
+					# Run the test
+					test_metrics, train_time, train_gflops, train_mem = tester.test()
+					logger.debug(f"test_metrics={test_metrics} time={train_time} gflops={train_gflops} mem={train_mem}")
+					
+					# Accomulate
+					total_test_metrics += test_metrics
+					model_id = report_model_names.index(model_name)
+					report_stats[model_id] += np.array([train_time, train_gflops, train_mem])
+
+					# Update the progress bar
+					c_current += 1
+					pbar.update(1)
+
+				# Calculate the stats (average over seeds)
+				test_metrics_mean = total_test_metrics / len(config.test_config["seeds"])
+				if len(report_params) < len(report_model_names):
+					report_params.append(count_parameters(model))
+				# Add to the current "line"
+				for m in test_metrics_mean:
+					report_line.append(f"{m:.3f}")
+			# Add line to the final report
+			report_data.append(report_line)
+	pbar.close()
+
+	# Calculate the  report stats
+	report_stats = report_stats.T
+	report_stats = report_stats / c_total
+
+	report_time 	= [f"{v:.3f}" for v in report_stats[0]]
+	report_gflops 	= [f"{v:.2f}" for v in report_stats[1]]
+	report_mem 		= [f"{v:.2f}" for v in report_stats[2]]
+	report_params 	= [f"{v:.2f}" for v in report_params  ]
+	# Obtain the columns&rows names
+	
+	report_columns = pd.MultiIndex.from_product([report_model_names, report_metric_names], names=["Model", "Metric"])
+	report_rows = pd.MultiIndex.from_product([report_dataset_names, config.test_config["contexts"]])
+	# Pack to the dataframe
+	df = pd.DataFrame(data = report_data, index = report_rows, columns = report_columns)
+
+	model_stats={
+		"Time (s/epoch)": report_time,
+		"Params count (K)": report_params,
+		"GFLOPs (f/batch)": report_gflops,
+		"Peak Memory (Gb)": report_mem
+	}
+
+	report = Report()
+	# Form and save the report
+	report.report(df, config.title, config.subtitle, model_stats)
+
+	return 0

modelmark/models/conv.py

@@ -0,0 +1,75 @@
+import torch
+import torch.nn as nn
+
+from modelmark.common.utils import load_config
+
+class ConvModel(nn.Module):
+	def __init__(self, input_dim : int, output_dim : int, context_size : int):
+		super().__init__()
+
+		# To avoid circular import, now it lives here
+		config = load_config()
+
+		self.input_size = input_dim
+		self.output_size = output_dim
+		self.context_size = context_size
+		
+		self.hidden_size = config.model_config["Conv"]["hidden_size"]
+		self.num_layers = config.model_config["Conv"]["num_layers"]
+		self.kernel_size = config.model_config["Conv"]["kernel_size"]
+		
+		layers = []
+
+		in_channels = self.input_size
+
+		for _ in range(self.num_layers):
+			layers.append(
+				nn.Conv1d(
+					in_channels=in_channels,
+					out_channels=self.hidden_size,
+					kernel_size=self.kernel_size,
+					padding=self.kernel_size // 2,
+				)
+			)
+			layers.append(nn.ReLU())
+
+			in_channels = self.hidden_size
+
+		self.conv = nn.Sequential(*layers)
+
+		# Collapse the sequence dimension to one feature vector
+		self.pool = nn.AdaptiveAvgPool1d(1)
+
+		# Produce exactly context_size * output_size values
+		self.fc = nn.Linear(
+			self.hidden_size,
+			context_size * self.output_size,
+		)
+
+	def forward(self, x):
+		# x: [batch, sequence_length, input_size]
+
+		# Conv1d expects:
+		# [batch, channels, sequence_length]
+		x = x.transpose(1, 2)
+
+		# [batch, hidden_size, sequence_length]
+		x = self.conv(x)
+
+		# [batch, hidden_size, 1]
+		x = self.pool(x)
+
+		# [batch, hidden_size]
+		x = x.squeeze(-1)
+
+		# [batch, context_size * output_size]
+		out = self.fc(x)
+
+		# [batch, context_size, output_size]
+		out = out.view(
+			x.size(0),
+			self.context_size,
+			self.output_size,
+		)
+
+		return out