diff --git a/.github/workflows/black.yml b/.github/workflows/black.yml
index f4196005..b48085c0 100644
--- a/.github/workflows/black.yml
+++ b/.github/workflows/black.yml
@@ -9,4 +9,4 @@ jobs:
- uses: actions/checkout@v2
- uses: psf/black@stable
with:
- src: "./simple_triton"
+ src: "./triteia"
diff --git a/.gitignore b/.gitignore
index f65d9bc0..14190d27 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,4 @@
+out
plot_out/
*.csv
*_results/
diff --git a/README.md b/README.md
index 858f3a6b..e35a447f 100644
--- a/README.md
+++ b/README.md
@@ -1,8 +1,8 @@
-# simple-triton
+# triteia
-simple-triton is a Python client for performing inference on the NVIDIA Triton Inference Server. It provides model deployment, configuration, and optimization capabilities for the TensorFlow, ONNX, and Python Triton backends directly from Python. This was developed to address limitations in the [PyTriton](https://github.com/triton-inference-server/pytriton) package that only suppports deployments with the Python backend where TensorRT, XLA, and mixed precision are not available.
+triteia is a Python client for performing inference on the NVIDIA Triton Inference Server. It provides model deployment, configuration, and optimization capabilities for the TensorFlow, ONNX, and Python Triton backends directly from Python. This was developed to address limitations in the [PyTriton](https://github.com/triton-inference-server/pytriton) package that only suppports deployments with the Python backend where TensorRT, XLA, and mixed precision are not available.
-
+
# User guide
@@ -22,25 +22,25 @@ simple-triton is a Python client for performing inference on the NVIDIA Triton I
## Quick start
-simple-triton requires `histomcs_stream` and `large_image` packages with the tiff reader
+triteia requires `histomcs_stream` and `large_image` packages with the tiff reader
```
-git clone https://github.com/PathologyDataScience/simple_triton.git
-pip install --editable ./simple_triton
+git clone https://github.com/PathologyDataScience/triteia.git
+pip install --editable ./triteia
```
-> `--editable` ensures that updates to the `simple_triton` package (after `git pull`) immediately takes effect.
+> `--editable` ensures that updates to the `triteia` package (after `git pull`) immediately takes effect.
-Or, you can try the Docker image. First, run `git clone` (as above) or make sure to do a git pull inside the "simple_triton" directory. Then:
+Or, you can try the Docker image. First, run `git clone` (as above) or make sure to do a git pull inside the "triteia" directory. Then:
```bash
# optional: download test data
python download_test_data.py
-# simple_triton_client will be the name of the Docker image
-docker build -f client.Dockerfile . -t simple_triton_client:latest --build-arg DOCKER_GROUP_ID=$(getent group docker | cut -d: -f3)
+# triteia will be the name of the Docker image
+docker build -f client.Dockerfile . -t triteia:latest --build-arg DOCKER_GROUP_ID=$(getent group docker | cut -d: -f3)
docker run \
--security-opt seccomp:unconfined --network=host \
--rm -it --shm-size=1g \
-v ${PWD}/test_data:/data:ro \
- --name tritonclient simple_triton_client:latest
+ --name tritonclient triteia:latest
```
> **_NOTE:_** `--network=` option allows the Docker image to access ports from other containers or the host. The default shared memory size for Docker containers is 64MB; use `--shm-size=` to increase it if you need to process large whole-slide images. The `--security-opt seccomp:unconfined` option may be needed on larger machines to enable [OpenBLAS](https://www.openblas.net/) threading support. The `--rm` option removes the container after it stops, so be cautious if you need persistent data.
@@ -64,13 +64,13 @@ docker run \
-v ${PWD}/examples:/examples:rw \
--user $UID --rm -it \
--shm-size=1g \
- --name tritonclient simple_triton_client:latest \
+ --name tritonclient triteia:latest \
bash -c "jupyter-lab --notebook-dir /examples/ --no-browser"
```
### Running the Triton container
-simple-triton is tested with [Triton version 25.02](https://github.com/triton-inference-server/server/releases/tag/v2.55.0).
+triteia is tested with [Triton version 25.02](https://github.com/triton-inference-server/server/releases/tag/v2.55.0).
Support for Tensorflow is deprecated in later versions, but other model backends (like PyTorch) should still work.
We recommend starting from the repository’s root directory (the same as the directory containing this README.md file). You can run using the `./launch_server.sh` script (which also has some command-line options) or the command below:
@@ -153,7 +153,7 @@ python feature_extraction.py ~/inputs.tsv ~/ EfficientNetV2S.tensorflow -s
```
## Model wrappers
-simple-triton contains wrappers for serving popular digital pathology models, including CONCH, UNI, Prov-GigaPath, hibou-L, Phikon, Virchow, and Virchow2 on the [Python backend](https://docs.nvidia.com/deeplearning/triton-inference-server/user-guide/docs/python_backend/README.html). All models are served using mixed precision.
+triteia contains wrappers for serving popular digital pathology models, including CONCH, UNI, Prov-GigaPath, hibou-L, Phikon, Virchow, and Virchow2 on the [Python backend](https://docs.nvidia.com/deeplearning/triton-inference-server/user-guide/docs/python_backend/README.html). All models are served using mixed precision.
| Model | Input | Output | Size |
|---|---|---|---|
@@ -195,14 +195,14 @@ docker run \
```
## Model configuration
-`simple_triton.config` includes model configuration classes that implement backend-specific configuration options. These classes enable configuration of batching behavior, specification of model input/output shapes and types, and backend optimizations. Refer to the Triton documentation on [model configuration](https://docs.nvidia.com/deeplearning/triton-inference-server/user-guide/docs/user_guide/model_configuration.html#model-configuration) and [optimization](https://docs.nvidia.com/deeplearning/triton-inference-server/user-guide/docs/user_guide/optimization.html) for further details.
+`triteia.config` includes model configuration classes that implement backend-specific configuration options. These classes enable configuration of batching behavior, specification of model input/output shapes and types, and backend optimizations. Refer to the Triton documentation on [model configuration](https://docs.nvidia.com/deeplearning/triton-inference-server/user-guide/docs/user_guide/model_configuration.html#model-configuration) and [optimization](https://docs.nvidia.com/deeplearning/triton-inference-server/user-guide/docs/user_guide/optimization.html) for further details.
Configuration classes like `PythonConfiguration` and `TensorflowConfiguration` take additional data classes as inputs that configure batching and caching behavior, hardware resources, and model input/output signatures.
When Triton is launched with `--strict-model-config=false`, the server automatically configures basic information such as input/output signatures, and the configuration can omit them.
```python
-from simple_triton.config import TensorflowConfig
-from simple_triton.model import TritonModel
+from triteia.config import TensorflowConfig
+from triteia.model import TritonModel
name = "mymodel.tensorflow"
config = TensorflowConfig(name, max_batch_size=64)
model = TritonModel(name, "localhost:8001")
@@ -211,7 +211,7 @@ model.load(config=config.json())
Alternatively, model inputs and output signatures can be defined using the `ModelInput` and `ModelOutput` classes
```python
-from simple_triton.config import ModelInput
+from triteia.config import ModelInput
input = [ModelInput(name="input_0", shape=[224, 224, 3], dtype=np.float32, optional=False)]
config = TensorflowConfig(name, max_batch_size=64, input=input)
```
@@ -220,14 +220,14 @@ Variable-sized input dimensions can be indicated using a value of -1.
The `InstanceGroup` class configures the use of CPU or GPU resources and the number of model instances hosted on each GPU.
```python
-from simple_triton.config import InstanceGroup
+from triteia.config import InstanceGroup
instances = InstanceGroup(count=2, kind="gpu", gpus=[0,1,2,3])
config = TensorflowConfig(name=name, instance_group=instances)
```
`TensorflowOptimization` can be used with `TensorflowMixedPrecision`, `TensorflowXla`, and `TensorRt` to activate automatic mixed precision, XLA compilation, or TensorRT optimization.
```python
-from simple_triton.config import TensorflowMixedPrecision, TensorflowXla, TensorflowOptimization
+from triteia.config import TensorflowMixedPrecision, TensorflowXla, TensorflowOptimization
amp = TensorflowMixedPrecision()
xla = TensorflowXla(level=2)
optimizer = TensorflowOptimization(amp=amp, xla=xla)
@@ -253,11 +253,11 @@ File-based configuration is useful for distributing models. When configuring and
The `TritonModel` class can be used to load/unload models, retrieve model configurations or metadata, or check whether a model is idle or loaded. A model is defined by a model name and a server URL.
```python
-from simple_triton.model import TritonModel
+from triteia.model import TritonModel
model = TritonModel("EfficientNetV2S.tensorflow", "localhost:8001")
```
-When unloading a model, simple-triton will check that the model is idle.
+When unloading a model, triteia will check that the model is idle.
```python
# load model with auto-generated configuration
# block and timeout after 1 second
@@ -306,7 +306,7 @@ pooch.retrieve(
### Using standalone docker container
To test using the Docker container, launch and build the client as follows:
```
-docker build -f client.Dockerfile . -t simple_triton_client:latest --build-arg DOCKER_GROUP_ID=$(getent group docker | cut -d: -f3)
+docker build -f client.Dockerfile . -t triteia:latest --build-arg DOCKER_GROUP_ID=$(getent group docker | cut -d: -f3)
./launch_test_container.sh
```
You can now run tests inside the container using `pytest tests` inside the container.
@@ -320,8 +320,11 @@ Be aware that extra performance is not guaranteed.
PyTorch offers many optimizations that may not be available in the ONNX or TRT backends.
## Paper results
-To reproduce the (TBD) paper: `OUTPUT_DIR="./results" ./benchmarking/paper_benchmarks.sh $OUTPUT_DIR`
-Results can be inspected either as tensorboards: `tensorboard --logdir=...`, or as figures:
+To reproduce the (TBD) paper
+1. start the NVIDIA triton server with all GPUs and models available: `./launch_server.sh --num-gpus --http-port 7984 --grpc-port 7985 --metrics-port 7986`
+2. `OUTPUT_DIR="./results" ./benchmarking/paper_benchmarks.sh $OUTPUT_DIR`
+3. Inspect results: `tensorboard --logdir=...`.
+To view results as figures::
```bash
# convert TensorBoard to CSV
./benchmarking/tensorboard_to_csv.py
@@ -330,5 +333,5 @@ Results can be inspected either as tensorboards: `tensorboard --logdir=...`, or
```
To do the benchmarks using Docker, use the `benchmark_client.Dockerfile` in the benchmarking directory.
-It is identical to the `client.Dockerfile` in this directory, except it has access to CUDA so it can automatically start and stop Triton with GPUs.
-To build it from the git root directory: `docker build -f benchmarking/benchmark_client.Dockerfile . -t simple_triton_client:benchmark`
+It is identical to the `client.Dockerfile` in this directory, except it has access to CUDA so it can export GPU metrics and automatically start and stop Triton with GPUs.
+To build it from the git root directory: `docker build -f benchmarking/benchmark_client.Dockerfile . -t triteia:benchmark`
diff --git a/benchmarking/ConvNeXtXLarge_amp_Batch64_GPU8_iter5_BatchTest.sh b/benchmarking/ConvNeXtXLarge_amp_Batch64_GPU8_iter5_BatchTest.sh
index 073c2abe..c9e01a62 100755
--- a/benchmarking/ConvNeXtXLarge_amp_Batch64_GPU8_iter5_BatchTest.sh
+++ b/benchmarking/ConvNeXtXLarge_amp_Batch64_GPU8_iter5_BatchTest.sh
@@ -22,7 +22,7 @@ while [ $gpu_num -ne 8 ]
do
maxbatchsize=$(($maxbatchsize+32))
echo "gpu_num:$gpu_num Workers: $limit maxbatchsize: $maxbatchsize"
- python /tf/notebooks/simple_triton/benchmarking/benchmark_interface.py --limit 1 --gpu-num $gpu_num --fileoutput $filename --iterations 5 --use-amp --precision "FP16" --maxbatchsize $maxbatchsize --model-name "ConvNeXtXLarge"
+ python /tf/notebooks/triteia/benchmarking/benchmark_interface.py --limit 1 --gpu-num $gpu_num --fileoutput $filename --iterations 5 --use-amp --precision "FP16" --maxbatchsize $maxbatchsize --model-name "ConvNeXtXLarge"
echo "==============================================================="
done
diff --git a/benchmarking/benchmark_client.Dockerfile b/benchmarking/benchmark_client.Dockerfile
index c1cfec0b..1e905933 100644
--- a/benchmarking/benchmark_client.Dockerfile
+++ b/benchmarking/benchmark_client.Dockerfile
@@ -1,5 +1,6 @@
# for the paper, this docker image is built with:
-# docker build -f client.Dockerfile . -t simple_triton_client:benchmark --build-arg DOCKER_GROUP_ID=$(getent group docker | cut -d: -f3) --build-arg UID=$(id -u) --build-arg GID=$(id -g) --build-arg USERNAME=$USER
+# this file is intended to be similar to the client.Dockerfile, except with GPUs available to the client (to read performance metrics, etc).
+# docker build -f client.Dockerfile . -t triteia:benchmark --build-arg DOCKER_GROUP_ID=$(getent group docker | cut -d: -f3) --build-arg UID=$(id -u) --build-arg GID=$(id -g) --build-arg USERNAME=$USER
FROM python:3.10-slim AS build-image
ARG USERNAME=myuser
@@ -14,9 +15,9 @@ RUN echo 'APT::Install-Suggests "0";' >> /etc/apt/apt.conf.d/00-docker && \
apt clean && \
rm -rf /var/lib/apt/lists/*
-# install simple-triton
-WORKDIR /home/$USERNAME/code/simple_triton
-COPY simple_triton/ simple_triton
+# install triteia
+WORKDIR /home/$USERNAME/code/triteia
+COPY triteia/ triteia
COPY pyproject.toml .
# comment out scm (i.e. git) line in pyproject.toml
RUN sed -i 's/.*\[tool.setuptools_scm\]/#&/g' pyproject.toml
@@ -44,8 +45,8 @@ RUN echo 'APT::Install-Suggests "0";' >> /etc/apt/apt.conf.d/00-docker && \
rm -rf /var/lib/apt/lists/*
USER $USERNAME
-WORKDIR /home/$USERNAME/simple_triton
-COPY --chown=$USERNAME:$USERNAME simple_triton/ simple_triton
+WORKDIR /home/$USERNAME/triteia
+COPY --chown=$USERNAME:$USERNAME triteia/ triteia
COPY --chown=$USERNAME:$USERNAME pyproject.toml .
@@ -75,8 +76,8 @@ RUN echo 'APT::Install-Suggests "0";' >> /etc/apt/apt.conf.d/00-docker && \
rm -rf /var/lib/apt/lists/*
RUN curl -fsSL https://get.docker.com | sh
# for jupyter notebooks as non-root
-RUN mkdir --mode a+rxw /.local /.jupyter /.cache /models/ /.config
-RUN chown $USERNAME:$USERNAME /home/$USERNAME/simple_triton/
+RUN mkdir --mode a+rxw /.local /.jupyter /.cache /.config
+RUN chown $USERNAME:$USERNAME /home/$USERNAME/triteia/
USER $USERNAME
COPY --chown=$USERNAME:$USERNAME README.md pyproject.toml ./
diff --git a/benchmarking/benchmark_interface.py b/benchmarking/benchmark_interface.py
index 161a08ec..d1c9ab19 100644
--- a/benchmarking/benchmark_interface.py
+++ b/benchmarking/benchmark_interface.py
@@ -9,11 +9,11 @@
import tensorflow as tf
from large_image.cache_util import cachesClear
-from simple_triton.config import *
-from simple_triton.feature_extraction import study, inference
-from simple_triton.model import TritonModel
-from simple_triton.tile_iterators import TiffPrefetch
-from simple_triton.utils import analyze
+from triteia.config import *
+from triteia.feature_extraction import study, inference
+from triteia.model import TritonModel
+from triteia.tile_iterators import TiffPrefetch
+from triteia.utils import analyze
class Benchmark:
@@ -23,6 +23,9 @@ class Benchmark:
"""
def __init__(self, args_dict):
+ self.times = None
+ self.tile_info = None
+ self.features = None
self.args_dict = args_dict
def create_hs_study(self, wsi_path, mask_path):
@@ -126,7 +129,7 @@ def callback(user_data, result, error):
)
# warm up Model
print("Warmup Model")
- (self.features, self.tile_info, self.times, self.failed,) = inference(
+ self.features, self.tile_info, self.times = inference(
iterator,
model_name,
url=self.args_dict["url"],
@@ -152,7 +155,7 @@ def callback(user_data, result, error):
)
# start timer
start = time.time()
- (self.features, self.tile_info, self.times, self.failed,) = inference(
+ self.features, self.tile_info, self.times = inference(
iterator,
model_name,
url=self.args_dict["url"],
@@ -212,8 +215,8 @@ def gpu_mem_clear(self):
def install():
"""Install dependencies for running benchmarking interface tool"""
# install large_image with tile sources as prereq, check feature_extraction.ipynb in examples directory.
- # install simple_triton
- subprocess.check_call([sys.executable, "-m", "pip", "install", f"../simple_triton"])
+ # install triteia
+ subprocess.check_call([sys.executable, "-m", "pip", "install", f"../triteia"])
subprocess.check_call([sys.executable, "-m", "pip", "install", "ray"])
subprocess.check_call([sys.executable, "-m", "pip", "install", "pyarrow"])
diff --git a/benchmarking/cpu_vs_gpu_plot.py b/benchmarking/cpu_vs_gpu_plot.py
new file mode 100644
index 00000000..8be38a30
--- /dev/null
+++ b/benchmarking/cpu_vs_gpu_plot.py
@@ -0,0 +1,171 @@
+"""
+Plot regular inference, inference without IO and Multiuser performance with and without TRT
+"""
+import string
+
+import matplotlib.pyplot as plt
+import pandas as pd
+import seaborn as sns
+from matplotlib.lines import Line2D
+
+# Extracted from bs128.
+# Not reading from tensorboards for this plot just to save some time and complexity
+data = """
+model_name,gpu_count,gpu_throughput,cpu_throughput
+ResNet-50,1,1475,1285
+ResNet-50,2,2606,1749
+ResNet-50,4,2814,1889
+ResNet-50,6,2947,1841
+ResNet-50,8,2830,1831
+UNI,1,0675,0570
+UNI,2,1310,1022
+UNI,4,2279,1539
+UNI,6,2425,1829
+UNI,8,2392,1853
+Prov-GigaPath,1,0212,0202
+Prov-GigaPath,2,0429,0407
+Prov-GigaPath,4,0844,0792
+Prov-GigaPath,6,1252,1092
+Prov-GigaPath,8,1602,1292
+"""
+
+
+# Colors using colorblind palette
+palette = sns.color_palette("colorblind", n_colors=3)
+colors = {
+ # f"CPU-Based": palette[0],
+ # f"GPU-Based": palette[1],
+ f"CPU-Based": "#1f77b4",
+ f"GPU-Based": "#ff7f0e",
+}
+
+# Parse the data
+df = pd.read_csv(pd.io.common.StringIO(data))
+
+# Define the order of models: Prov-Gigapath, UNI, then ResNet-50
+model_order = ["Prov-GigaPath", "UNI", "ResNet-50"]
+
+# Create figure with 3 subplots (one for each model)
+fig, axes = plt.subplots(1, 3, figsize=(15, 5), gridspec_kw={"hspace": 0.5})
+fig.suptitle(
+ "CPU- vs GPU-based Preprocessing Throughput Comparison",
+ fontsize=16,
+ fontweight="bold",
+)
+fig.subplots_adjust(top=0.82) # Add vertical space after the suptitle
+
+# GPU counts for x-axis
+gpu_counts = [1, 2, 4, 6, 8]
+x_positions = range(len(gpu_counts))
+bar_width = 0.35
+
+label_idx = 0
+
+# Plot each model in its own subplot
+for idx, model in enumerate(model_order):
+ ax = axes[idx]
+
+ # Filter data for this model
+ model_data = df[df["model_name"] == model].sort_values("gpu_count")
+
+ # Extract values for plotting
+ gpu_throughput = model_data["gpu_throughput"].values
+ cpu_throughput = model_data["cpu_throughput"].values
+
+ # Create bars for GPU and CPU
+ x_pos = list(x_positions)
+ ax.bar(
+ [x - bar_width / 2 for x in x_pos],
+ cpu_throughput,
+ bar_width,
+ label="CPU-Based",
+ color=colors["CPU-Based"],
+ edgecolor="black",
+ linewidth=0.5,
+ )
+ ax.bar(
+ [x + bar_width / 2 for x in x_pos],
+ gpu_throughput,
+ bar_width,
+ label="GPU-Based",
+ color=colors["GPU-Based"],
+ edgecolor="black",
+ linewidth=0.5,
+ )
+
+ # Set labels and title
+ ax.set_xlabel("GPUs", fontsize=16)
+ ax.set_title(model, fontsize=12, fontweight="bold")
+ ax.set_xticks(x_pos)
+ ax.set_xticklabels(gpu_counts, fontsize=14)
+ ax.grid(axis="y", alpha=0.3)
+
+ letter = string.ascii_lowercase[label_idx]
+ ax.text(
+ 0.09,
+ 1.02,
+ f"{letter})",
+ transform=ax.transAxes,
+ fontsize=16,
+ fontweight="bold",
+ va="bottom",
+ ha="right",
+ )
+ label_idx += 1
+
+ # Only show y-axis label and ticks on the first subplot
+ if idx == 0:
+ ax.set_ylabel("tiles/s", fontsize=16)
+ else:
+ ax.set_yticklabels([])
+
+# Set the same y-axis limits for all subplots
+all_throughputs = list(df["gpu_throughput"]) + list(df["cpu_throughput"])
+y_max = max(all_throughputs) * 1.1 # Add 10% margin
+for ax in axes:
+ ax.set_ylim(0, y_max)
+
+# Create a single legend below the figure
+handles = [
+ plt.Rectangle((0, 0), 1, 1, fc=colors["CPU-Based"], label="CPU-Based"),
+ plt.Rectangle((0, 0), 1, 1, fc=colors["GPU-Based"], label="GPU-Based"),
+]
+fig.legend(
+ handles=handles,
+ loc="upper center",
+ bbox_to_anchor=(0.5, -0.02),
+ ncol=2,
+ frameon=True,
+ fontsize=18,
+)
+
+plt.tight_layout()
+
+figure_dst = "plot_out/cpu-vs-gpu.png"
+plt.savefig(figure_dst, dpi=300, bbox_inches="tight")
+
+# Also save as SVG
+figure_dst_svg = "plot_out/cpu-vs-gpu.svg"
+plt.savefig(figure_dst_svg, format="svg", bbox_inches="tight")
+
+# plt.show()
+plt.close()
+print(f"Saved figure to {figure_dst}")
+print(f"Saved figure to {figure_dst_svg}")
+
+print("\n=== Percentage Difference Analysis (GPU vs CPU) ===\n")
+for model in model_order:
+ model_data = df[df["model_name"] == model].sort_values("gpu_count")
+ print(f"{model}:")
+ for _, row in model_data.iterrows():
+ gpu_count = row["gpu_count"]
+ gpu_throughput = row["gpu_throughput"]
+ cpu_throughput = row["cpu_throughput"]
+
+ # Calculate percentage difference: (GPU - CPU) / CPU * 100
+ pct_diff = ((gpu_throughput - cpu_throughput) / cpu_throughput) * 100
+
+ print(
+ f" GPU Count {int(gpu_count)}: {pct_diff:+.1f}% (GPU: {gpu_throughput}, CPU: {cpu_throughput})"
+ )
+ print()
diff --git a/benchmarking/gpu_vs_cpu_comparison.sh b/benchmarking/gpu_vs_cpu_comparison.sh
new file mode 100755
index 00000000..2965e373
--- /dev/null
+++ b/benchmarking/gpu_vs_cpu_comparison.sh
@@ -0,0 +1,23 @@
+#!/usr/bin/env bash
+# See README.md
+#
+# This file creates number from CPU-only based preprocessing
+# the model definitions are not included in the repository.
+# to create them, duplicate the gigapath, uni and python into "_old" folders. Change the device from "self.device" to CPU where appropriate.
+
+OUTPUT_DIR="${1:-/results/}"
+CLEAR_CACHE_REMOTELY="${2:-false}"
+
+for modelname in uni_old.python gigapath_old.python resnet_old.python
+do
+ ./benchmarking/triton_batchsize_benchmark.sh /data/5 $modelname "${OUTPUT_DIR}" "${CLEAR_CACHE_REMOTELY}" false
+done
+
+# for modelname in uni_old.python gigapath_old.python resnet_old.python
+# do
+# ./benchmarking/triton_batchsize_benchmark.sh /data/7 $modelname "${OUTPUT_DIR}" "${CLEAR_CACHE_REMOTELY}" 1
+# done
+#just in case
+#docker container stop tritonserver_$USER || true
+
+echo "All benchmarks done"
diff --git a/benchmarking/paper_benchmarks.sh b/benchmarking/paper_benchmarks.sh
index 0309b70b..a6b0d900 100755
--- a/benchmarking/paper_benchmarks.sh
+++ b/benchmarking/paper_benchmarks.sh
@@ -49,14 +49,14 @@ done
#just in case
docker container stop tritonserver_$USER || true
-
+
for modelname in resnet50 uni gigapath resnet50_trt_uint8 #gigapath_trt_uint8
do
- ./benchmarking/triton_batchsize_benchmark.sh /data/5 $modelname "${OUTPUT_DIR}" "${CLEAR_CACHE_REMOTELY}" 1 1,8 256
+ ./benchmarking/triton_batchsize_benchmark.sh /data/5 $modelname "${OUTPUT_DIR}" "${CLEAR_CACHE_REMOTELY}" 1 #1,8 256
done
-#just in case
+# just in case
docker container stop tritonserver_$USER || true
-
+#
for modelname in resnet50 uni gigapath
do
./benchmarking/pytorch_batchsize_benchmark.sh /data/5 $modelname "${OUTPUT_DIR}" "${CLEAR_CACHE_REMOTELY}" false
@@ -64,7 +64,7 @@ done
for modelname in resnet50 uni gigapath
do
- ./benchmarking/pytorch_batchsize_benchmark.sh /data/5 $modelname "${OUTPUT_DIR}" "${CLEAR_CACHE_REMOTELY}" 1 1,8 256
+ ./benchmarking/pytorch_batchsize_benchmark.sh /data/5 $modelname "${OUTPUT_DIR}" "${CLEAR_CACHE_REMOTELY}" 1 #1,8 256
done
for modelname in resnet50 uni gigapath resnet50_trt_uint8 #gigapath_trt_uint8
@@ -75,11 +75,15 @@ done
docker container stop tritonserver_$USER || true
-for modelname in resnet50 uni gigapath
-do
- ./benchmarking/triton_limit_benchmark.sh /data/5 $modelname "${OUTPUT_DIR}" "${CLEAR_CACHE_REMOTELY}" 1 8 16,20
-done
+# to test the impact of "--limit" parameter. Did not find anything interesting (general rule of thumb: limit should about 2x number of GPUs)
+# for modelname in resnet50 uni gigapath
+# do
+# ./benchmarking/triton_limit_benchmark.sh /data/5 $modelname "${OUTPUT_DIR}" "${CLEAR_CACHE_REMOTELY}" 1 8 16,20
+# done
# just in case
-docker container stop tritonserver_$USER || true
+# docker container stop tritonserver_$USER || true
+#
+
+./benchmarking/gpu_vs_cpu_comparison.sh
echo "All benchmarks done"
diff --git a/benchmarking/plot_with_gpu_scaling.py b/benchmarking/plot_with_gpu_scaling.py
new file mode 100644
index 00000000..fb7229d9
--- /dev/null
+++ b/benchmarking/plot_with_gpu_scaling.py
@@ -0,0 +1,750 @@
+#!/usr/bin/env bash
+# -*- coding: utf-8 -*-
+"""
+Read tensorboarddata from benchmarks and write them to CSV files.
+This way it is easy to modify the plot code without having to parse all the tensorboards.
+"""
+
+import argparse
+import glob
+import logging
+import os
+import re
+import sys
+
+import matplotlib.pyplot as plt
+import matplotlib.lines as mlines
+import pandas as pd
+
+GPU_SCALING_CSV = "gpu_scaling_ramin.csv"
+
+
+def parse_args():
+ parser = argparse.ArgumentParser(description="Plot tensorboard data")
+ parser.add_argument(
+ "--tensorboard-dir",
+ type=str,
+ nargs="+",
+ help="One or more directories containing tensorboard runs (or parent directories of runs)",
+ required=True,
+ default=["."],
+ )
+
+ parser.add_argument(
+ "--output-dir",
+ type=str,
+ help="output directory for plot files",
+ default="./plot_out",
+ )
+
+ args = parser.parse_args()
+
+ if not os.path.exists(args.output_dir):
+ os.makedirs(args.output_dir)
+
+ for d in args.tensorboard_dir:
+ if not os.path.exists(d) or not os.path.isdir(d):
+ raise ValueError(
+ f"Tensorboard directory '{d}' does not exist or is not a directory"
+ )
+
+ return args
+
+
+def extract_tensorboard_data(logdir: str, plot_key: str) -> pd.DataFrame:
+ """
+ Read scalar summaries for `plot_key` from a TensorBoard logdir.
+
+ Returns a dataframe with columns: step, wall_time, value.
+
+ Note: TensorFlow is imported lazily here so that if gpu_scaling.csv exists,
+ the program can start quickly without importing TensorFlow at all.
+ """
+ import tensorflow as tf
+
+ event_files = glob.glob(
+ os.path.join(logdir, "**", "events.out.tfevents.*"), recursive=True
+ )
+ if not event_files:
+ logging.warning("No event files found under %s", logdir)
+ return pd.DataFrame(columns=["step", "wall_time", "value"])
+
+ rows = []
+ for ef in sorted(event_files):
+ try:
+ for e in tf.compat.v1.train.summary_iterator(ef):
+ if not hasattr(e, "summary") or e.summary is None:
+ raise ValueError(
+ f"Event file {ef} does not contain a valid summary"
+ )
+
+ for v in e.summary.value:
+ if v.tag != plot_key:
+ continue
+ # Prefer simple_value when present
+ if hasattr(v, "simple_value"):
+ value = float(v.simple_value)
+ else:
+ # Some summaries store tensor; skip if we can't read it robustly here
+ continue
+ rows.append(
+ {
+ "step": int(getattr(e, "step", 0)),
+ "wall_time": float(getattr(e, "wall_time", 0.0)),
+ "value": value,
+ }
+ )
+ except Exception as dle:
+ logging.warning(
+ "Skipping event file %s due to data loss error: %s", ef, dle
+ )
+
+ if not rows:
+ return pd.DataFrame(columns=["step", "wall_time", "value"])
+
+ df = pd.DataFrame(rows).sort_values(["step", "wall_time"]).reset_index(drop=True)
+ return df
+
+
+def _candidate_run_dirs(tensorboard_dirs):
+ # Collect candidate run directories from the provided list (roots + direct children)
+ candidate_run_dirs = []
+ for root in tensorboard_dirs:
+ candidate_run_dirs.append(root)
+ for name in os.listdir(root):
+ p = os.path.join(root, name)
+ if os.path.isdir(p):
+ candidate_run_dirs.append(p)
+ return sorted(set(candidate_run_dirs))
+
+
+def build_gpu_scaling_dataframe(
+ tensorboard_dirs, plot_key: str, run_prefix="", run_suffix=""
+) -> pd.DataFrame:
+ """
+ Scan tensorboard directories and build a single dataframe of run-level summaries.
+
+ Output columns:
+ model, run, logdir, gpus, batch_size, plot_key, throughput_mean, throughput_last, gpu_util_percent_mean
+
+ If run_prefix is non-empty, the run dir must start with that prefix, e.g. "pytorch".
+ If run_suffix is non-empty, the run dir must end with that suffix, e.g. "inferenceonlytriton".
+ """
+ prefix = re.escape(run_prefix)
+ suffix = re.escape(run_suffix)
+
+ model_specs = [
+ {
+ "name": (
+ f"{prefix}Prov-GigaPath{suffix}"
+ if (run_prefix or run_suffix)
+ else "GigaPath"
+ ),
+ "run_dir_re": re.compile(rf"{prefix}gigapathgpu[1-8]bs[0-9]+{suffix}$"),
+ "parse_re": re.compile(
+ rf"{prefix}gigapathgpu(?P[1-8])bs(?P[0-9]+){suffix}$"
+ ),
+ },
+ {
+ "name": (f"{prefix}UNI{suffix}" if (run_prefix or run_suffix) else "UNI"),
+ "run_dir_re": re.compile(rf"{prefix}unigpu[1-8]bs[0-9]+{suffix}$"),
+ "parse_re": re.compile(
+ rf"{prefix}unigpu(?P[1-8])bs(?P[0-9]+){suffix}$"
+ ),
+ },
+ {
+ "name": (
+ f"{prefix}ResNet-50{suffix}"
+ if (run_prefix or run_suffix)
+ else "ResNet-50"
+ ),
+ "run_dir_re": re.compile(rf"{prefix}resnet50gpu[1-8]bs[0-9]+{suffix}$"),
+ "parse_re": re.compile(
+ rf"{prefix}resnet50gpu(?P[1-8])bs(?P[0-9]+){suffix}$"
+ ),
+ },
+ {
+ "name": (
+ f"{prefix}ResNet-50-trt{suffix}"
+ if (run_prefix or run_suffix)
+ else "ResNet-50-trt"
+ ),
+ "run_dir_re": re.compile(
+ rf"{prefix}resnet50_trt_uint8gpu[1-8]bs[0-9]+{suffix}$"
+ ),
+ "parse_re": re.compile(
+ rf"{prefix}resnet50_trt_uint8gpu(?P[1-8])bs(?P[0-9]+){suffix}$"
+ ),
+ # tritonresnet50_trt_uint8gpu2bs64
+ },
+ {
+ "name": (
+ f"{prefix}Prov-GigaPath-trt{suffix}"
+ if (run_prefix or run_suffix)
+ else "Prov-GigaPath-trt"
+ ),
+ "run_dir_re": re.compile(
+ rf"{prefix}gigapath_trt_uint8gpu[1-8]bs[0-9]+{suffix}$"
+ ),
+ "parse_re": re.compile(
+ rf"{prefix}gigapath_trt_uint8gpu(?P[1-8])bs(?P[0-9]+){suffix}$"
+ ),
+ },
+ ]
+
+ candidate_run_dirs = _candidate_run_dirs(tensorboard_dirs)
+
+ rows = []
+ for spec in model_specs:
+ run_dirs = []
+ for d in candidate_run_dirs:
+ base = os.path.basename(os.path.normpath(d))
+ if spec["run_dir_re"].match(base):
+ run_dirs.append(d)
+
+ if not run_dirs:
+ logging.warning(
+ "No tensorboard run directories matched for model '%s' (prefix=%r, suffix=%r)",
+ spec["name"],
+ run_prefix,
+ run_suffix,
+ )
+ continue
+ for d in run_dirs:
+ base = os.path.basename(os.path.normpath(d))
+ m = spec["parse_re"].match(base)
+ if not m:
+ continue
+
+ gpus = int(m.group("gpus"))
+ bs = int(m.group("bs"))
+
+ df = extract_tensorboard_data(d, plot_key=plot_key)
+ if df.empty:
+ logging.warning("No data for key '%s' in %s", plot_key, d)
+ continue
+
+ # Extract GPU utilization data — average across all GPUs in this run
+ gpu_util_means = []
+ for gpu_idx in range(gpus):
+ gpu_util_key = f"gpu_gpu_{gpu_idx}_util_percent"
+ gpu_util_df = extract_tensorboard_data(d, plot_key=gpu_util_key)
+ if not gpu_util_df.empty:
+ gpu_util_means.append(float(gpu_util_df["value"].mean()))
+ else:
+ logging.warning("No data for key '%s' in %s", gpu_util_key, d)
+
+ gpu_util_mean = (
+ float(sum(gpu_util_means) / len(gpu_util_means))
+ if gpu_util_means
+ else None
+ )
+
+ rows.append(
+ {
+ "model": spec["name"],
+ "run": base,
+ "logdir": d,
+ "gpus": gpus,
+ "batch_size": bs,
+ "plot_key": plot_key,
+ "throughput_mean": float(df["value"].mean()),
+ "throughput_last": float(df["value"].iloc[-1]),
+ "gpu_util_percent_mean": gpu_util_mean,
+ }
+ )
+
+ out = pd.DataFrame(
+ rows,
+ columns=[
+ "model",
+ "run",
+ "logdir",
+ "gpus",
+ "batch_size",
+ "plot_key",
+ "throughput_mean",
+ "throughput_last",
+ "gpu_util_percent_mean",
+ ],
+ )
+
+ if not out.empty:
+ out = out.sort_values(["model", "batch_size", "gpus", "run"]).reset_index(
+ drop=True
+ )
+ return out
+
+
+def load_or_build_gpu_scaling(
+ tensorboard_dirs,
+ csv_path: str,
+ plot_key: str,
+ run_prefix="triton",
+ run_suffix="",
+) -> pd.DataFrame:
+ """
+ If csv exists, load it and skip reading TensorBoards.
+ Otherwise, build dataframe from TensorBoards and write csv.
+ """
+ if os.path.exists(csv_path) and os.path.isfile(csv_path):
+ logging.info(
+ "Found %s; loading cached data (skipping TensorBoard read).", csv_path
+ )
+ df = pd.read_csv(csv_path)
+ return df
+
+ logging.info("%s not found; reading TensorBoards and creating cache.", csv_path)
+ df = build_gpu_scaling_dataframe(
+ tensorboard_dirs,
+ plot_key=plot_key,
+ run_prefix=run_prefix,
+ run_suffix=run_suffix,
+ )
+ if df.empty:
+ raise ValueError(
+ f"No matching runs with scalar data found for plot key '{plot_key}' when requesting {csv_path} from {tensorboard_dirs}."
+ )
+
+ df.to_csv(csv_path, index=False)
+ logging.info("Wrote %d rows to %s", len(df), csv_path)
+ return df
+
+
+def read_and_write_data(tensorboard_dirs, output_dir):
+ plot_key = "throughput_total_tiles_per_second"
+
+ for run_prefix in ["triton"]:
+ for run_suffix in ["", ": pre-loaded slides"]:
+ key = "inferenceonly" if run_suffix else ""
+ df = load_or_build_gpu_scaling(
+ tensorboard_dirs=tensorboard_dirs,
+ csv_path=os.path.join(
+ output_dir, f"{run_prefix}{key}{GPU_SCALING_CSV}"
+ ),
+ plot_key=plot_key,
+ run_prefix=run_prefix,
+ run_suffix=key,
+ )
+
+
+def plot_batch_size_throughput_gpu_util(
+ df: pd.DataFrame,
+ title: str = "Batch size vs Throughput & GPU Utilization",
+ output_dir: str = "./plot_out",
+ output_filename: str = "batch_size_throughput_gpu_util.png",
+ gpus_filter: int | None = None,
+):
+ """
+ For each model, plot a figure with batch size on the x-axis,
+ throughput (tiles/s) on the left y-axis, and GPU utilization (%)
+ on the right y-axis. Two curves per subplot: one for throughput
+ and one for GPU utilization.
+
+ Parameters
+ ----------
+ df : pd.DataFrame
+ Must contain columns: model, batch_size, throughput_mean, gpu_util_percent_mean.
+ Optionally 'gpus' to filter by GPU count.
+ title : str
+ Super-title for the figure.
+ output_dir : str
+ Directory to save the figure.
+ output_filename : str
+ File name for the saved figure.
+ gpus_filter : int or None
+ If set, only rows with this many GPUs are plotted.
+ """
+ df = df.copy()
+
+ # Clean model names (strip framework prefixes)
+ df["model"] = (
+ df["model"]
+ .str.replace("triton", "", regex=False)
+ .str.replace("pytorch", "", regex=False)
+ .str.replace("inferenceonly", "", regex=False)
+ )
+
+ for col in ["batch_size", "throughput_mean", "gpu_util_percent_mean"]:
+ df[col] = pd.to_numeric(df[col], errors="coerce")
+ df = df.dropna(subset=["batch_size", "throughput_mean", "gpu_util_percent_mean"])
+
+ if gpus_filter is not None:
+ df["gpus"] = pd.to_numeric(df["gpus"], errors="coerce")
+ df = df[df["gpus"] == gpus_filter]
+
+ if df.empty:
+ logging.warning("No data to plot for batch-size throughput/GPU-util figure.")
+ return
+
+ # Aggregate per (model, batch_size)
+ agg = (
+ df.groupby(["model", "batch_size"], as_index=False)
+ .agg(
+ throughput_mean=("throughput_mean", "mean"),
+ gpu_util_percent_mean=("gpu_util_percent_mean", "mean"),
+ )
+ .sort_values(["model", "batch_size"])
+ )
+
+ models = sorted(agg["model"].unique())
+ # move the last entry to the second position
+ models = [models[0], models[-1]] + models[1:-1]
+ n_models = len(models)
+
+ fig, axes = plt.subplots(1, n_models, figsize=(6 * n_models, 5), squeeze=False)
+
+ color_throughput = "#1f77b4"
+ color_gpu_util = "#ff7f0e"
+
+ # Compute shared y-limits across all models
+ throughput_max = agg["throughput_mean"].max()
+ throughput_ylim = (0, throughput_max * 1.1)
+ gpu_util_ylim = (0, 105)
+
+ all_ax_right = []
+
+ for idx, model in enumerate(models):
+ ax_left = axes[0, idx]
+ model_data = agg[agg["model"] == model].sort_values("batch_size")
+
+ batch_sizes = model_data["batch_size"].values
+ throughputs = model_data["throughput_mean"].values
+ gpu_utils = model_data["gpu_util_percent_mean"].values
+
+ # Use evenly-spaced integer positions so ticks are equally spaced
+ x_positions = list(range(len(batch_sizes)))
+
+ # Left y-axis: throughput
+ ax_left.plot(
+ x_positions,
+ throughputs,
+ marker="o",
+ color=color_throughput,
+ linewidth=2,
+ )
+ ax_left.set_xlabel("Batch Size", fontsize=13)
+ ax_left.set_xticks(x_positions)
+ ax_left.set_xticklabels([str(int(bs)) for bs in batch_sizes])
+ ax_left.set_title(model, fontsize=14, fontweight="bold")
+ ax_left.grid(True, alpha=0.3)
+ ax_left.set_ylim(throughput_ylim)
+
+ # Only show left y-axis label and tick labels on the first subplot
+ if idx == 0:
+ ax_left.set_ylabel("tiles / s", fontsize=13, color=color_throughput)
+ ax_left.tick_params(axis="y", labelcolor=color_throughput)
+ else:
+ ax_left.set_ylabel("")
+ ax_left.tick_params(axis="y", labelleft=False)
+
+ # Right y-axis: GPU utilization
+ ax_right = ax_left.twinx()
+ all_ax_right.append(ax_right)
+ ax_right.plot(
+ x_positions,
+ gpu_utils,
+ marker="s",
+ color=color_gpu_util,
+ linewidth=2,
+ linestyle="--",
+ )
+ ax_right.set_ylim(gpu_util_ylim)
+
+ # Only show right y-axis label and tick labels on the last subplot
+ if idx == n_models - 1:
+ ax_right.set_ylabel(
+ "GPU Utilization (%)", fontsize=13, color=color_gpu_util
+ )
+ ax_right.tick_params(axis="y", labelcolor=color_gpu_util)
+ else:
+ ax_right.set_ylabel("")
+ ax_right.tick_params(axis="y", labelright=False)
+
+ # Shared legend at the bottom of the figure
+ handle_throughput = mlines.Line2D(
+ [],
+ [],
+ color=color_throughput,
+ marker="o",
+ linewidth=2,
+ label="tiles/s",
+ )
+ handle_gpu_util = mlines.Line2D(
+ [],
+ [],
+ color=color_gpu_util,
+ marker="s",
+ linewidth=2,
+ linestyle="--",
+ label="GPU util %",
+ )
+ fig.legend(
+ handles=[handle_throughput, handle_gpu_util],
+ loc="upper center",
+ bbox_to_anchor=(0.5, -0.02),
+ ncol=4,
+ frameon=True,
+ fontsize=12,
+ )
+
+ fig.suptitle(title, fontsize=16, fontweight="bold")
+ plt.tight_layout(rect=[0, 0.03, 1, 0.95])
+
+ os.makedirs(output_dir, exist_ok=True)
+ plot_path = os.path.join(output_dir, output_filename)
+ plt.savefig(plot_path, dpi=300, bbox_inches="tight")
+ logging.info("Saved plot to %s", plot_path)
+ plt.close()
+
+
+STANDARD_BATCH_SIZES = sorted([32, 64, 128, 256])
+ALLOWED_GPU_TICKS = [1, 2, 4, 6, 8]
+
+
+def plot_gpu_scaling_throughput_by_batch_size(
+ df: pd.DataFrame,
+ title: str = "GPU scaling — Throughput by batch size",
+ output_dir: str = "./plot_out",
+ output_filename: str = "gpu_scaling_throughput_by_batchsize.png",
+):
+ """
+ One subplot per model.
+ X-axis: number of GPUs (evenly spaced ticks for 1, 2, 4, 6, 8).
+ Left y-axis (shared): throughput (tiles/s) — solid lines, one per batch size.
+ Right y-axis (shared): GPU utilization (%) — dashed lines, same color per batch size.
+ A single shared legend at the bottom.
+ """
+ import seaborn as sns
+
+ df = df.copy()
+
+ # Clean model names
+ df["model"] = (
+ df["model"]
+ .str.replace("triton", "", regex=False)
+ .str.replace("pytorch", "", regex=False)
+ .str.replace("inferenceonly", "", regex=False)
+ )
+
+ for col in ["gpus", "batch_size", "throughput_mean", "gpu_util_percent_mean"]:
+ df[col] = pd.to_numeric(df[col], errors="coerce")
+ df = df.dropna(subset=["gpus", "batch_size", "throughput_mean"])
+
+ # Aggregate per (model, gpus, batch_size)
+ agg_cols = {"throughput_mean": ("throughput_mean", "mean")}
+ if df["gpu_util_percent_mean"].notna().any():
+ agg_cols["gpu_util_percent_mean"] = ("gpu_util_percent_mean", "mean")
+ agg = (
+ df.groupby(["model", "gpus", "batch_size"], as_index=False)
+ .agg(**agg_cols)
+ .sort_values(["model", "batch_size", "gpus"])
+ )
+ has_gpu_util = "gpu_util_percent_mean" in agg.columns
+
+ available_batch_sizes = sorted(
+ [bs for bs in STANDARD_BATCH_SIZES if bs in agg["batch_size"].unique()]
+ )
+ if not available_batch_sizes:
+ logging.warning("No standard batch sizes found for GPU-scaling plot.")
+ return
+
+ models = sorted(agg["model"].unique())
+ n_models = len(models)
+ if n_models == 0:
+ logging.warning("No models found for GPU-scaling plot.")
+ return
+
+ # Shared y-limits
+ y_max = agg["throughput_mean"].max() * 1.1
+ gpu_util_ylim = (0, 105)
+
+ # Evenly-spaced x positions for GPU ticks
+ gpu_x = {gpu: i for i, gpu in enumerate(ALLOWED_GPU_TICKS)}
+ x_positions = list(range(len(ALLOWED_GPU_TICKS)))
+
+ # Colors & markers per batch size
+ palette = sns.color_palette("colorblind", n_colors=len(available_batch_sizes))
+ bs_colors = dict(zip(available_batch_sizes, palette))
+ markers_solid = ["o", "s", "D", "^"]
+ markers_open = ["v", "P", "X", "*"]
+ bs_markers_throughput = dict(
+ zip(available_batch_sizes, markers_solid[: len(available_batch_sizes)])
+ )
+ bs_markers_gpu_util = dict(
+ zip(available_batch_sizes, markers_open[: len(available_batch_sizes)])
+ )
+
+ fig, axes = plt.subplots(1, n_models, figsize=(6 * n_models, 5), squeeze=False)
+
+ for idx, model in enumerate(models):
+ ax_left = axes[0, idx]
+ model_data = agg[agg["model"] == model]
+
+ # --- Left y-axis: throughput (solid lines) ---
+ for bs in available_batch_sizes:
+ bs_data = model_data[model_data["batch_size"] == bs].sort_values("gpus")
+ xs = [gpu_x[g] for g in bs_data["gpus"].values if g in gpu_x]
+ ys = bs_data.loc[
+ bs_data["gpus"].isin(gpu_x.keys()), "throughput_mean"
+ ].values
+ ax_left.plot(
+ xs,
+ ys,
+ marker=bs_markers_throughput[bs],
+ color=bs_colors[bs],
+ linewidth=2,
+ )
+
+ ax_left.set_xlabel("GPUs", fontsize=13)
+ ax_left.set_xticks(x_positions)
+ ax_left.set_xticklabels([str(g) for g in ALLOWED_GPU_TICKS], fontsize=12)
+ ax_left.set_title(model, fontsize=14, fontweight="bold")
+ ax_left.grid(True, alpha=0.3)
+ ax_left.set_ylim(0, y_max)
+
+ # Only show left y-axis label/ticks on the leftmost subplot
+ if idx == 0:
+ ax_left.set_ylabel("tiles / s", fontsize=13)
+ else:
+ ax_left.set_ylabel("")
+ ax_left.tick_params(axis="y", labelleft=False)
+
+ # --- Right y-axis: GPU utilization (dashed lines) ---
+ if has_gpu_util:
+ ax_right = ax_left.twinx()
+ for bs in available_batch_sizes:
+ bs_data = model_data[model_data["batch_size"] == bs].sort_values("gpus")
+ xs = [gpu_x[g] for g in bs_data["gpus"].values if g in gpu_x]
+ ys = bs_data.loc[
+ bs_data["gpus"].isin(gpu_x.keys()), "gpu_util_percent_mean"
+ ].values
+ ax_right.plot(
+ xs,
+ ys,
+ marker=bs_markers_gpu_util[bs],
+ color=bs_colors[bs],
+ linewidth=2,
+ linestyle="--",
+ alpha=0.7,
+ )
+ ax_right.set_ylim(gpu_util_ylim)
+
+ # Only show right y-axis label/ticks on the rightmost subplot
+ if idx == n_models - 1:
+ ax_right.set_ylabel("GPU Utilization (%)", fontsize=13)
+ else:
+ ax_right.set_ylabel("")
+ ax_right.tick_params(axis="y", labelright=False)
+
+ # --- Single shared legend at the bottom ---
+ handles = []
+ for bs in available_batch_sizes:
+ # Throughput handle (solid)
+ handles.append(
+ mlines.Line2D(
+ [],
+ [],
+ color=bs_colors[bs],
+ marker=bs_markers_throughput[bs],
+ linewidth=2,
+ label=f"bs={int(bs)} tiles/s",
+ )
+ )
+ if has_gpu_util:
+ for bs in available_batch_sizes:
+ # GPU util handle (dashed)
+ handles.append(
+ mlines.Line2D(
+ [],
+ [],
+ color=bs_colors[bs],
+ marker=bs_markers_gpu_util[bs],
+ linewidth=2,
+ linestyle="--",
+ alpha=0.7,
+ label=f"bs={int(bs)} GPU util %",
+ )
+ )
+
+ fig.legend(
+ handles=handles,
+ loc="upper center",
+ bbox_to_anchor=(0.5, -0.02),
+ ncol=len(available_batch_sizes),
+ frameon=True,
+ fontsize=11,
+ )
+
+ fig.suptitle(title, fontsize=16, fontweight="bold")
+ plt.tight_layout(rect=[0, 0.03, 1, 0.95])
+
+ os.makedirs(output_dir, exist_ok=True)
+ plot_path = os.path.join(output_dir, output_filename)
+ plt.savefig(plot_path, dpi=300, bbox_inches="tight")
+ logging.info("Saved plot to %s", plot_path)
+ plt.close()
+
+
+def main():
+ args = parse_args()
+ read_and_write_data(args.tensorboard_dir, args.output_dir)
+
+ # --- Generate batch-size vs throughput & GPU utilization plots ---
+ for run_prefix in ["triton"]:
+ for key_label, key in [("", ""), ("inferenceonly", "inferenceonly")]:
+ csv_path = os.path.join(
+ args.output_dir, f"{run_prefix}{key}{GPU_SCALING_CSV}"
+ )
+ if not os.path.exists(csv_path):
+ logging.info("CSV not found, skipping: %s", csv_path)
+ continue
+
+ df = pd.read_csv(csv_path)
+ suffix_label = (
+ " (inference only)" if key == "inferenceonly" else " (with I/O)"
+ )
+
+ # If multiple GPU counts exist, plot one figure per GPU count
+ # if "gpus" in df.columns:
+ # for gpus in sorted(df["gpus"].unique()):
+ # plot_batch_size_throughput_gpu_util(
+ # df,
+ # title=f"{run_prefix.capitalize()}{suffix_label} — {int(gpus)} GPU(s)",
+ # output_dir=args.output_dir,
+ # output_filename=f"{run_prefix}{key}_batch_throughput_gpuutil_{int(gpus)}gpu.png",
+ # gpus_filter=int(gpus),
+ # )
+ # else:
+ # plot_batch_size_throughput_gpu_util(
+ # df,
+ # title=f"{run_prefix.capitalize()}{suffix_label}",
+ # output_dir=args.output_dir,
+ # output_filename=f"{run_prefix}{key}_batch_throughput_gpuutil.png",
+ # )
+
+ # --- Generate GPU-scaling plots ---
+ for run_prefix in ["triton"]:
+ for key_label, key in [("", ""), ("inferenceonly", "inferenceonly")]:
+ csv_path = os.path.join(
+ args.output_dir, f"{run_prefix}{key}{GPU_SCALING_CSV}"
+ )
+ if not os.path.exists(csv_path):
+ logging.info("CSV not found, skipping: %s", csv_path)
+ continue
+
+ df = pd.read_csv(csv_path)
+ suffix_label = (
+ " (inference only)" if key == "inferenceonly" else " (with I/O)"
+ )
+ plot_gpu_scaling_throughput_by_batch_size(
+ df,
+ title=f"{run_prefix.capitalize()}{suffix_label}",
+ output_dir=args.output_dir,
+ output_filename=f"{run_prefix}{key}_gpu_scaling.png",
+ )
+
+
+if __name__ == "__main__":
+ logging.basicConfig(stream=sys.stdout, level=logging.INFO)
+ main()
diff --git a/benchmarking/pytorch_benchmark.py b/benchmarking/pytorch_benchmark.py
index 2bb36200..12a0bd67 100644
--- a/benchmarking/pytorch_benchmark.py
+++ b/benchmarking/pytorch_benchmark.py
@@ -23,9 +23,9 @@
ResNetModel,
)
-path.append(os.path.join(os.path.dirname(__file__), "../simple_triton"))
-from simple_triton.feature_extraction import study
-from simple_triton.tile_iterators import TiffPrefetch
+path.append(os.path.join(os.path.dirname(__file__), "../triteia"))
+from triteia.feature_extraction import study
+from triteia.tile_iterators import TiffPrefetch
from util import (
clear_cache,
convert_seconds_to_hms,
@@ -34,7 +34,7 @@
write_energy_stats,
)
-from simple_triton.utils import init_tb_writer, track_method
+from triteia.utils import init_tb_writer, track_method
def normalize_image(in_0, device, mean, std):
diff --git a/benchmarking/run_docker.sh b/benchmarking/run_docker.sh
index 25dae74a..698f38c3 100755
--- a/benchmarking/run_docker.sh
+++ b/benchmarking/run_docker.sh
@@ -1,4 +1,10 @@
#!/usr/bin/env bash
+#
+# This file launches docker and mounts some directories
+# if you have issues with missing files in the server:
+# The model directory will be also be mounted in the server, which is launched from THIS CLIENT FILE
+# in other words, a recursive mount: docker run -v /dir:/dir .... docker run -v /dir:/dir
+# keep in mind that recursive mounts with docker will always rely on HOST paths, not from within a container
# make sure we are in root directory of the github repo
cd "$(dirname "$0")"/..
@@ -16,12 +22,13 @@ docker run \
--network=host \
--init \
--shm-size=20g \
- -v "$PWD/launch_server.sh":/home/$USER/simple_triton/launch_server.sh \
- -v "$PWD/benchmarking":/home/$USER/simple_triton/benchmarking \
+ -v "$PWD/launch_server.sh":/home/$USER/triteia/launch_server.sh \
+ -v "$PWD/benchmarking":/home/$USER/triteia/benchmarking \
+ -v "$PWD/models/":/home/$USER/triteia/models/ \
-v /home/aza4423/BENCHMARK_DATA/:/data:ro \
- -v "$PWD/test_data":/home/$USER/simple_triton/test_data/ \
- -v /data/anders_aza4423/simple_triton_results/:/results \
+ -v "$PWD/test_data":/home/$USER/triteia/test_data/ \
+ -v /data/anders_aza4423/triteia_results/:/results \
-v /var/run/docker.sock:/var/run/docker.sock \
--rm \
- --name tritonclient_test_$USER \
- -it simple_triton_client:benchmark
\ No newline at end of file
+ --name triteia_test_$USER \
+ -it triteia:benchmark
diff --git a/benchmarking/tensorboard_csv_to_plot.py b/benchmarking/tensorboard_csv_to_plot.py
index ae1ee689..f9038724 100644
--- a/benchmarking/tensorboard_csv_to_plot.py
+++ b/benchmarking/tensorboard_csv_to_plot.py
@@ -1,7 +1,7 @@
#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
-Plot data from benchmarks of simple_triton.
+Plot data from benchmarks of triteia.
First run the `tensorboard_parser.py`
This code is mostly generated by AI agents.
You may consider splitting into smaller files to avoid wasting tokens.
@@ -11,8 +11,11 @@
import argparse
import logging
import os
+import string
+
import sys
+import matplotlib.axes
import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
@@ -30,6 +33,18 @@
ALLOWED_GPU_TICKS = [1, 2, 4, 6, 8]
STANDARD_BATCH_SIZES = sorted([32, 64, 128, 256])
+AVERAGE_TILE_READ_SPEED_B32 = 1661
+AVERAGE_TILE_READ_SPEED_B64 = 2649
+AVERAGE_TILE_READ_SPEED_B128 = 3843
+AVERAGE_TILE_READ_SPEED_B256 = 3845
+
+AVERAGE_TILE_READ_SPEED = {
+ 32: AVERAGE_TILE_READ_SPEED_B32,
+ 64: AVERAGE_TILE_READ_SPEED_B64,
+ 128: AVERAGE_TILE_READ_SPEED_B128,
+ 256: AVERAGE_TILE_READ_SPEED_B256,
+}
+
# ---------------------------------------------------------------------------
# Common helpers
@@ -123,6 +138,7 @@ def _create_model_subplots(
n_rows,
n_cols,
figsize=(col_width * n_cols, row_height * n_rows),
+ gridspec_kw={"hspace": 0.5},
**subplot_kw,
)
# Normalise axes to always be indexable
@@ -169,6 +185,29 @@ def _plot_overlapping_bars(
ax.bar(gpu_idx, height, bar_width, **kw)
+def _plot_grouped_bars(
+ ax, gpu_groups, category_values, color_map, bar_width=None, alpha_map=None
+):
+ """Draw side-by-side (grouped) bars on *ax*, one sub-bar per category value."""
+ n_cats = len(category_values)
+ if bar_width is None:
+ bar_width = 0.8 / max(n_cats, 1)
+
+ for gpu_idx, gpu in enumerate(ALLOWED_GPU_TICKS):
+ for cat_idx, cat in enumerate(category_values):
+ offset = (cat_idx - (n_cats - 1) / 2) * bar_width
+ height = gpu_groups[gpu][cat]
+ if height > 0:
+ kw = dict(
+ color=color_map.get(cat, "gray"),
+ edgecolor="black",
+ linewidth=0.5,
+ )
+ if alpha_map:
+ kw["alpha"] = alpha_map.get(cat, 0.5)
+ ax.bar(gpu_idx + offset, height, bar_width, **kw)
+
+
def _format_gpu_axis(ax, ax_idx, *, ylabel="tiles / s", ylim_max=None, fontsize=14):
"""Apply common GPU-axis formatting."""
x = range(len(ALLOWED_GPU_TICKS))
@@ -193,13 +232,25 @@ def _add_bottom_legend(fig, handles, ncol):
)
+def _save_figure_formats(fig, plot_path, **savefig_kwargs):
+ """Save figure as PNG plus a matching SVG file."""
+ root, ext = os.path.splitext(plot_path)
+ svg_path = f"{root}.svg" if ext else f"{plot_path}.svg"
+ output_paths = [plot_path]
+ if os.path.abspath(svg_path) != os.path.abspath(plot_path):
+ output_paths.append(svg_path)
+
+ for output_path in output_paths:
+ fig.savefig(output_path, **savefig_kwargs)
+ logging.info("Saved plot to %s", output_path)
+
+
def _save_and_close(fig, plot_path, suptitle=None, rect=None):
- """Set suptitle, tight_layout, save, log, close."""
+ """Set suptitle, tight_layout, save as PNG/SVG, log, close."""
if suptitle:
fig.suptitle(suptitle, fontsize=16, fontweight="bold", y=0.98)
plt.tight_layout(rect=rect or [0, 0.03, 1, 0.96])
- plt.savefig(plot_path, bbox_inches="tight", dpi=300)
- logging.info("Saved plot to %s", plot_path)
+ _save_figure_formats(fig, plot_path, bbox_inches="tight", dpi=300)
plt.close()
@@ -324,6 +375,7 @@ def plot_gpu_scaling_stacked_bar_from_df(
)
)
+ label_idx = 0
for ax_idx, model in enumerate(models):
ax = axes[ax_idx]
model_data = agg_df[agg_df["model"] == model]
@@ -334,6 +386,20 @@ def plot_gpu_scaling_stacked_bar_from_df(
_format_gpu_axis(ax, ax_idx, ylim_max=ylim_max)
ax.set_title(model)
+ # Add subplot letter label (a, b, c, ...)
+ letter = string.ascii_lowercase[label_idx]
+ label_idx += 1
+ ax.text(
+ 0.09,
+ 1.02,
+ f"{letter})",
+ transform=ax.transAxes,
+ fontsize=16,
+ fontweight="bold",
+ va="bottom",
+ ha="right",
+ )
+
handles = _make_rect_handles(available_batch_sizes, batch_colors, label_fmt="bs={}")
_add_bottom_legend(fig, handles, ncol=len(available_batch_sizes))
@@ -398,6 +464,7 @@ def plot_multiuser_single_concurrency_combined(
n_models = len(models)
x = range(len(ALLOWED_GPU_TICKS))
+ label_idx = 0
for model_idx, model in enumerate(models):
offset = (model_idx - (n_models - 1) / 2) * bar_width
heights = []
@@ -416,6 +483,14 @@ def plot_multiuser_single_concurrency_combined(
edgecolor="black",
linewidth=0.5,
)
+ # letter = string.ascii_lowercase[label_idx]
+ # ax.text(
+ # 0.09, 1.02, f"{letter})",
+ # transform=ax.transAxes,
+ # fontsize=16, fontweight="bold",
+ # va="bottom", ha="right",
+ # )
+ label_idx += 1
ax.set_xlabel("GPUs", fontsize=16)
ax.set_ylabel(y_axis_label, fontsize=16)
@@ -429,8 +504,7 @@ def plot_multiuser_single_concurrency_combined(
plt.tight_layout()
filename = f"multiuser_c{concurrency}_gpu_scaling_combined.png"
plot_path = os.path.join(output_dir, f"{output_prefix}{filename}")
- plt.savefig(plot_path, dpi=300, bbox_inches="tight")
- logging.info("Saved plot to %s", plot_path)
+ _save_figure_formats(fig, plot_path, dpi=300, bbox_inches="tight")
plt.close()
@@ -467,6 +541,7 @@ def plot_inferenceonly_triton_pytorch(
}
bar_width = 0.35
+ label_idx = 0
for ax_idx, model in enumerate(models):
ax = axes[ax_idx]
model_data = agg_df[agg_df["model"] == model]
@@ -520,6 +595,18 @@ def plot_inferenceonly_triton_pytorch(
_format_gpu_axis(ax, ax_idx, fontsize=16)
ax.set_title(model, fontsize=14, fontweight="bold")
ax.set_ylim(0, ylim)
+ letter = string.ascii_lowercase[label_idx]
+ ax.text(
+ 0.09,
+ 1.02,
+ f"{letter})",
+ transform=ax.transAxes,
+ fontsize=16,
+ fontweight="bold",
+ va="bottom",
+ ha="right",
+ )
+ label_idx += 1
# Create legend with raw line indicator
handles, labels = axes[0].get_legend_handles_labels()
@@ -618,6 +705,7 @@ def plot_combined_triton_pytorch_raw(
gridspec_kw={"hspace": 0.5},
)
+ label_idx = 0
for row_idx, fw_name in enumerate(row_labels):
agg_df = agg_frames[fw_name]
for col_idx, model in enumerate(models):
@@ -640,6 +728,20 @@ def plot_combined_triton_pytorch_raw(
ax.grid(True, alpha=0.3, axis="y")
ax.set_ylim(0, ylim_max)
+ # Add subplot letter label (a, b, c, ...)
+ letter = string.ascii_lowercase[label_idx]
+ ax.text(
+ 0.09,
+ 1.02,
+ f"{letter})",
+ transform=ax.transAxes,
+ fontsize=16,
+ fontweight="bold",
+ va="bottom",
+ ha="right",
+ )
+ label_idx += 1
+
# Move first row (Triton) slightly lower
offset_amount = (
0.04 # Adjust this value to control how much lower the first row moves
@@ -685,6 +787,245 @@ def plot_combined_triton_pytorch_raw(
)
+def plot_combined_triton_pytorch_raw_grouped(
+ df_triton: pd.DataFrame,
+ df_pytorch: pd.DataFrame,
+ output_dir: str = "./plot_out",
+ plot_key: str = "throughput_total_tiles_per_second",
+ ylim_max: int = 3500,
+):
+ """Same as plot_combined_triton_pytorch_raw but with side-by-side bars."""
+ agg_frames = {}
+ for fw_name, df in {"Triton": df_triton, "PyTorch": df_pytorch}.items():
+ agg = _prepare_agg_frame(df, plot_key)
+ if agg is None or agg.empty:
+ logging.warning("No data for %s with plot_key=%s", fw_name, plot_key)
+ return
+ agg_frames[fw_name] = agg
+
+ all_models = set()
+ for agg_df in agg_frames.values():
+ all_models.update(agg_df["model"].unique())
+ models = [m for m in PREFERRED_MODELS if m in all_models]
+
+ all_bs = set()
+ for agg_df in agg_frames.values():
+ all_bs.update(agg_df["batch_size"].unique())
+ available_batch_sizes = [bs for bs in STANDARD_BATCH_SIZES if bs in all_bs]
+ if not available_batch_sizes:
+ logging.warning("No standard batch sizes found for combined plot")
+ return
+
+ batch_colors = dict(
+ zip(
+ available_batch_sizes,
+ sns.color_palette("colorblind", n_colors=len(available_batch_sizes)),
+ )
+ )
+
+ import string
+ from matplotlib.lines import Line2D
+
+ row_labels = ["Triton", "PyTorch"]
+ n_rows = len(row_labels)
+ n_cols = len(models)
+
+ # Compute grouped-bar width (must match _plot_grouped_bars default)
+ n_cats = len(available_batch_sizes)
+ bar_width = 0.8 / max(n_cats, 1)
+
+ fig, axes = _create_model_subplots(
+ models,
+ n_rows=n_rows,
+ col_width=5,
+ row_height=5,
+ sharey=True,
+ sharex=True,
+ gridspec_kw={"hspace": 0.5},
+ )
+
+ label_idx = 0
+ for row_idx, fw_name in enumerate(row_labels):
+ agg_df = agg_frames[fw_name]
+ for col_idx, model in enumerate(models):
+ ax: matplotlib.axes.Axes = (
+ axes[row_idx, col_idx] if n_rows > 1 else axes[col_idx]
+ )
+ model_data = agg_df[agg_df["model"] == model]
+ gpu_groups = _build_gpu_groups(
+ model_data, available_batch_sizes, "throughput_mean", "batch_size"
+ )
+ _plot_grouped_bars(ax, gpu_groups, available_batch_sizes, batch_colors)
+
+ if row_idx == n_rows - 1:
+ ax.set_xlabel("GPUs", fontsize=14)
+ x = range(len(ALLOWED_GPU_TICKS))
+ ax.set_xticks(list(x))
+ ax.set_xticklabels([str(g) for g in ALLOWED_GPU_TICKS], fontsize=14)
+ if col_idx == 0:
+ ax.set_ylabel("tiles / s", fontsize=14)
+ if row_idx == 0:
+ ax.set_title(model, fontsize=14, fontweight="bold")
+ ax.grid(True, alpha=0.3, axis="y")
+ ax.set_ylim(0, ylim_max)
+
+ # Add subplot letter label (a, b, c, ...)
+ letter = string.ascii_lowercase[label_idx]
+ ax.text(
+ 0.09,
+ 1.02,
+ f"{letter})",
+ transform=ax.transAxes,
+ fontsize=16,
+ fontweight="bold",
+ va="bottom",
+ ha="right",
+ )
+ label_idx += 1
+
+ # Move first row (Triton) slightly lower
+ offset_amount = (
+ 0.04 # Adjust this value to control how much lower the first row moves
+ )
+ for col_idx in range(n_cols):
+ ax = axes[0, col_idx]
+ pos = ax.get_position()
+ new_pos = [pos.x0, pos.y0 - offset_amount, pos.width, pos.height]
+ ax.set_position(new_pos)
+
+ # Add framework labels above each row
+ for row_idx, fw_name in enumerate(row_labels):
+ ax_first = axes[row_idx, 0] if n_rows > 1 else axes[0]
+ pos = ax_first.get_position()
+ fig.text(
+ 0.1,
+ pos.y1 + 0.05,
+ f"{fw_name}:",
+ ha="center",
+ va="bottom",
+ fontsize=16,
+ fontweight="bold",
+ )
+
+ row_labels = ["Triton", "PyTorch"]
+ n_rows = len(row_labels)
+ n_cols = len(models)
+
+ fig, axes = _create_model_subplots(
+ models,
+ n_rows=n_rows,
+ col_width=5,
+ row_height=5,
+ sharey=True,
+ sharex=True,
+ gridspec_kw={"hspace": 0.5},
+ )
+
+ label_idx = 0
+ for row_idx, fw_name in enumerate(row_labels):
+ agg_df = agg_frames[fw_name]
+ for col_idx, model in enumerate(models):
+ ax = axes[row_idx, col_idx] if n_rows > 1 else axes[col_idx]
+ model_data = agg_df[agg_df["model"] == model]
+ gpu_groups = _build_gpu_groups(
+ model_data, available_batch_sizes, "throughput_mean", "batch_size"
+ )
+ _plot_grouped_bars(ax, gpu_groups, available_batch_sizes, batch_colors)
+
+ if row_idx == n_rows - 1:
+ ax.set_xlabel("GPUs", fontsize=14)
+ x = range(len(ALLOWED_GPU_TICKS))
+ ax.set_xticks(list(x))
+ ax.set_xticklabels([str(g) for g in ALLOWED_GPU_TICKS], fontsize=14)
+ if col_idx == 0:
+ ax.set_ylabel("tiles / s", fontsize=14)
+ if row_idx == 0:
+ ax.set_title(model, fontsize=14, fontweight="bold")
+ ax.grid(True, alpha=0.3, axis="y")
+ ax.set_ylim(0, ylim_max)
+
+ # Draw average tile read speed as markers on the y-axis (first column only)
+ if col_idx >= 0:
+ for bs in available_batch_sizes:
+ speed = AVERAGE_TILE_READ_SPEED[bs]
+ # ax.plot(
+ # [0], speed,
+ # marker='>',
+ # markersize=10,
+ # color=batch_colors[bs],
+ # markeredgecolor='black',
+ # markeredgewidth=1,
+ # linestyle='none',
+ # zorder=5,
+ # )
+ ax.axhline(
+ y=speed,
+ color=batch_colors[bs],
+ linestyle="-",
+ linewidth=2.5,
+ zorder=3,
+ alpha=0.4,
+ )
+
+ # Add subplot letter label (a, b, c, ...)
+ letter = string.ascii_lowercase[label_idx]
+ ax.text(
+ 0.09,
+ 1.02,
+ f"{letter})",
+ transform=ax.transAxes,
+ fontsize=16,
+ fontweight="bold",
+ va="bottom",
+ ha="right",
+ )
+ label_idx += 1
+
+ # Move first row (Triton) slightly lower
+ offset_amount = (
+ 0.04 # Adjust this value to control how much lower the first row moves
+ )
+ for col_idx in range(n_cols):
+ ax = axes[0, col_idx]
+ pos = ax.get_position()
+ new_pos = [pos.x0, pos.y0 - offset_amount, pos.width, pos.height]
+ ax.set_position(new_pos)
+
+ # Add framework labels above each row
+ for row_idx, fw_name in enumerate(row_labels):
+ ax_first = axes[row_idx, 0] if n_rows > 1 else axes[0]
+ pos = ax_first.get_position()
+ fig.text(
+ 0.1,
+ pos.y1 + 0.05,
+ f"{fw_name}:",
+ ha="center",
+ va="bottom",
+ fontsize=16,
+ fontweight="bold",
+ )
+
+ handles = _make_rect_handles(available_batch_sizes, batch_colors, label_fmt="bs={}")
+ fig.legend(
+ handles=handles,
+ loc="upper center",
+ bbox_to_anchor=(0.5, 0.03),
+ ncol=len(available_batch_sizes),
+ frameon=True,
+ fontsize=18,
+ )
+
+ plot_path = os.path.join(
+ output_dir, "combined_triton_pytorch_throughput_batchsize_grouped.png"
+ )
+ _save_and_close(
+ fig,
+ plot_path,
+ suptitle="Batch size (bs) throughput scaling using Triton and PyTorch",
+ rect=[0.05, 0.03, 1, 0.96],
+ )
+
+
def main():
args = parse_args()
plot_key = "throughput_total_tiles_per_second"
@@ -726,11 +1067,15 @@ def main():
df_triton["framework"] = "Triton_raw"
df_pytorch["framework"] = "PyTorch_raw"
- plot_combined_triton_pytorch_raw(
- df_triton, df_pytorch, args.output_dir, ylim_max=3700
- )
+ # plot_combined_triton_pytorch_raw(
+ # df_triton, df_pytorch, args.output_dir, ylim_max=3700
+ # )
- # Create comparison plot for batch size 256
+ plot_combined_triton_pytorch_raw_grouped(
+ df_triton, df_pytorch, args.output_dir, ylim_max=4000
+ )
+ #
+ # # Create comparison plot for batch size 256
df_triton_inferenceonly["framework"] = "Triton"
df_pytorch_inferenceonly["framework"] = "PyTorch"
df_combined = pd.concat(
@@ -741,11 +1086,11 @@ def main():
df_combined, batch_size=256, output_dir=args.output_dir, output_prefix=""
)
- csv_path = os.path.join(args.output_dir, LIMIT_SCALING_CSV)
- limit_df = pd.read_csv(csv_path)
- plot_limit_scaling_from_df(limit_df, plot_key=plot_key, output_dir=args.output_dir)
+ # csv_path = os.path.join(args.output_dir, LIMIT_SCALING_CSV)
+ # limit_df = pd.read_csv(csv_path)
+ # plot_limit_scaling_from_df(limit_df, plot_key=plot_key, output_dir=args.output_dir)
- # Multiuser aggregate CSV + plot
+ # # Multiuser aggregate CSV + plot
csv_path = os.path.join(args.output_dir, MULTIUSER_SCALING_CSV)
if not os.path.exists(csv_path) or not os.path.isfile(csv_path):
logging.info("Not found %s; skipping Multiuser aggregate CSV + plot.", csv_path)
diff --git a/benchmarking/tensorboard_csv_to_plots_trt.py b/benchmarking/tensorboard_csv_to_plots_trt.py
index 9097ee39..1bdd3e71 100644
--- a/benchmarking/tensorboard_csv_to_plots_trt.py
+++ b/benchmarking/tensorboard_csv_to_plots_trt.py
@@ -1,6 +1,8 @@
"""
Plot regular inference, inference without IO and Multiuser performance with and without TRT
"""
+import string
+
import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
@@ -53,6 +55,7 @@
bar_width = 0.25
x_positions = {1: 0, 8: 1}
+label_idx = 0
for ax_idx, model_cfg in enumerate(models_config):
ax = axes[ax_idx]
@@ -159,6 +162,20 @@
linewidth=2.5,
)
+ # Add subplot letter label (a, b, c, ...)
+ letter = string.ascii_lowercase[label_idx]
+ ax.text(
+ 0.09,
+ 1.02,
+ f"{letter})",
+ transform=ax.transAxes,
+ fontsize=16,
+ fontweight="bold",
+ va="bottom",
+ ha="right",
+ )
+ label_idx += 1
+
ax.set_xlabel("GPUs", fontsize=14)
if ax_idx == 0:
ax.set_ylabel("tiles / second", fontsize=14)
@@ -194,6 +211,8 @@
fontsize=11,
)
-plt.savefig("plot_out/trt_comparison.png", dpi=300, bbox_inches="tight")
+figure_dst = "plot_out/trt_comparison.png"
+plt.savefig(figure_dst, dpi=300, bbox_inches="tight")
# plt.show()
plt.close()
+print(f"Saved figure to {figure_dst}")
diff --git a/benchmarking/tensorboard_to_csv.py b/benchmarking/tensorboard_to_csv.py
index bf3e9abe..f85e92d9 100644
--- a/benchmarking/tensorboard_to_csv.py
+++ b/benchmarking/tensorboard_to_csv.py
@@ -209,10 +209,11 @@ def build_gpu_scaling_dataframe(
if not run_dirs:
logging.warning(
- "No tensorboard run directories matched for model '%s' (prefix=%r, suffix=%r)",
+ "No tensorboard run directories matched for model '%s' (prefix=%r, suffix=%r) (regex: %r)",
spec["name"],
run_prefix,
run_suffix,
+ spec["run_dir_re"],
)
continue
for d in run_dirs:
@@ -386,7 +387,7 @@ def load_or_build_gpu_scaling(
)
if df.empty:
raise ValueError(
- f"No matching runs with scalar data found for plot key '{plot_key}'"
+ f"No matching runs with scalar data found for plot key '{plot_key}' with prefix '{run_prefix}' and suffix '{run_suffix}' in {tensorboard_dirs}"
)
df.to_csv(csv_path, index=False)
@@ -596,32 +597,31 @@ def read_and_write_data(tensorboard_dirs, output_dir):
plot_key = "throughput_total_tiles_per_second"
for run_prefix in ["triton", "pytorch"]:
- for run_suffix in ["", ": pre-loaded slides"]:
- key = "inferenceonly" if run_suffix else ""
+ for run_suffix in ["", "inferenceonly"]:
df = load_or_build_gpu_scaling(
tensorboard_dirs=tensorboard_dirs,
csv_path=os.path.join(
- output_dir, f"{run_prefix}{key}{GPU_SCALING_CSV}"
+ output_dir, f"{run_prefix}{run_suffix}{GPU_SCALING_CSV}"
),
plot_key=plot_key,
run_prefix=run_prefix,
- run_suffix=key,
+ run_suffix=run_suffix,
)
df = load_or_build_gpu_scaling(
tensorboard_dirs=tensorboard_dirs,
csv_path=os.path.join(
- output_dir, f"{run_prefix}{key}{LATENCY_GPU_SCALING_CSV}"
+ output_dir, f"{run_prefix}{run_suffix}{LATENCY_GPU_SCALING_CSV}"
),
plot_key="latency_mean_ms",
run_prefix=run_prefix,
- run_suffix=key,
+ run_suffix=run_suffix,
)
- limit_df = load_or_build_limit_scaling(
- tensorboard_dirs=tensorboard_dirs,
- csv_path=os.path.join(output_dir, LIMIT_SCALING_CSV),
- plot_key=plot_key,
- )
+ # limit_df = load_or_build_limit_scaling(
+ # tensorboard_dirs=tensorboard_dirs,
+ # csv_path=os.path.join(output_dir, LIMIT_SCALING_CSV),
+ # plot_key=plot_key,
+ # )
# Multiuser aggregate CSV + plot
multiuser_df = load_or_build_multiuser_scaling(
diff --git a/benchmarking/tileiterator-benchmark.py b/benchmarking/tileiterator-benchmark.py
index 734e3e09..248db951 100644
--- a/benchmarking/tileiterator-benchmark.py
+++ b/benchmarking/tileiterator-benchmark.py
@@ -11,7 +11,7 @@
to get estimated number of raw disk content for a given .svs:
(adjust the 'sed 1d;3d' and filename)
- tiffdump -m 100000 ~/simple_triton/test_data/wsi/TCGA-AN-A0G0-01Z-00-DX1.svs | grep TileByteCounts | sed '1d;3d' | awk '{ for (i=6; i<=NF; i++) sum += $i } END { print sum }'
+ tiffdump -m 100000 ~/triteia/test_data/wsi/TCGA-AN-A0G0-01Z-00-DX1.svs | grep TileByteCounts | sed '1d;3d' | awk '{ for (i=6; i<=NF; i++) sum += $i } END { print sum }'
"""
import csv
@@ -21,16 +21,16 @@
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime
-# ensure we're loading the simple_triton in this directory, not installed in path
+# ensure we're loading the triteia in this directory, not installed in path
from sys import path
from time import perf_counter
import numpy as np
from matplotlib import pyplot as plt
-path.append(os.path.join(os.path.dirname(__file__), "../simple_triton"))
-from simple_triton.feature_extraction import study
-from simple_triton.tile_iterators import TiffPrefetch
+path.append(os.path.join(os.path.dirname(__file__), "../triteia"))
+from triteia.feature_extraction import study
+from triteia.tile_iterators import TiffPrefetch
from util import parse_args, clear_cache
diff --git a/benchmarking/triton_batchsize_benchmark.sh b/benchmarking/triton_batchsize_benchmark.sh
index 25dc6767..00c07f58 100755
--- a/benchmarking/triton_batchsize_benchmark.sh
+++ b/benchmarking/triton_batchsize_benchmark.sh
@@ -3,7 +3,7 @@
wsi_path="${1:?Error: WSI path must be provided as \$1}"
modelname="${2:?Error: Model name must be provided as \$2}"
output_dir="${3:?Error: Output directory must be provided as \$3}"
-CLEAR_CACHE_REMOTELY="${4:-true}"
+CLEAR_CACHE_REMOTELY="${4:-false}"
inference_only="${5:-false}"
gpus="${6:-1,2,4,6,8}"
batch_sizes="${7:-32,64,128,256}"
@@ -18,7 +18,7 @@ fi
clear_cache() {
if [[ "${CLEAR_CACHE_REMOTELY,,}" == "true" ]]; then
echo "Clearing cache remotely..."
- curl localhost:7987/run
+ curl localhost:7987/run || exit 1
else
echo "Skipping remote cache clear (CLEAR_CACHE_REMOTELY=false)"
fi
@@ -55,18 +55,22 @@ do
fi
test -d "${output_dir}/${tb_name}" && continue
set -xe
- if [ ! "$(docker ps -q -f name=tritonserver_$USER)" ]; then
- ./launch_server.sh --num-gpus $gpu --start-gpu-id 0 ${modelname} --detached 1 --http-port 7984 --grpc-port 7985 --metrics-port 7986
- sleep 120
- fi
+ # The below code is useful for automatic stop and start of the server, but
+ # with "--manual-preload", all you need it to run the server with all GPUs
+ # and models available
+ # if [ ! "$(docker ps -q -f name=tritonserver_$USER)"
+ # ]; then
+ # ./launch_server.sh --num-gpus $gpu --start-gpu-id 0 ${modelname} --detached 1 --http-port 7984 --grpc-port 7985 --metrics-port 7986
+ # sleep 120
+ # fi
clear_cache
- eval "$cmd --batch-size ${bs} --tensorboard-name ${tb_name} --gpus ${gpu_string}"
+ eval "$cmd --batch-size ${bs} --tensorboard-name ${tb_name} --gpus ${gpu_string} --manual-preload"
set +xe
done
- if [ "$(docker ps -q -f name=tritonserver_$USER)" ]
- then
- docker container stop tritonserver_$USER
- sleep 10 # seems to be necessary - docker container stop does not properly clean up right away
- fi
+ # if [ "$(docker ps -q -f name=tritonserver_$USER)" ]
+ # then
+ # docker container stop tritonserver_$USER
+ # sleep 10 # seems to be necessary - docker container stop does not properly clean up right away
+ # fi
done
diff --git a/benchmarking/triton_benchmark.py b/benchmarking/triton_benchmark.py
index ef696d38..2bddd425 100644
--- a/benchmarking/triton_benchmark.py
+++ b/benchmarking/triton_benchmark.py
@@ -10,7 +10,7 @@
from contextlib import ExitStack
from pprint import pprint
-# ensure we're loading the simple_triton in this directory, not installed in path
+# ensure we're loading the triteia in this directory, not installed in path
from sys import path
from time import perf_counter
@@ -19,9 +19,9 @@
import torch
from tensorboardX import GlobalSummaryWriter
-# Make sure we're testing the local simple_triton
-path.append(os.path.join(os.path.dirname(__file__), "../simple_triton"))
-from config import PythonConfig, InstanceGroup
+# Make sure we're testing the local triteia
+path.append(os.path.join(os.path.dirname(__file__), "../triteia"))
+from config import PythonConfig, InstanceGroup, TensorRTConfig
from model import TritonModel
from feature_extraction import study, inference, inference_job, initialize_clients
from tile_iterators import TiffPrefetch
@@ -106,11 +106,22 @@ def main():
args = parse_args()
if not args.preload:
- config = PythonConfig(
- args.model_name,
- args.max_batch_size,
- instance_group=InstanceGroup(count=args.instance_group),
- )
+ if "trt" in args.model_name:
+ config = TensorRTConfig(
+ args.model_name,
+ args.max_batch_size,
+ instance_group=InstanceGroup(
+ count=args.instance_group, kind="gpu", gpus=args.gpus
+ ),
+ )
+ else:
+ config = PythonConfig(
+ args.model_name,
+ args.max_batch_size,
+ instance_group=InstanceGroup(
+ count=args.instance_group, kind="gpu", gpus=args.gpus
+ ),
+ )
model = TritonModel(args.model_name, args.url)
model.load(config=config.json())
assert model.is_loaded()
@@ -207,7 +218,7 @@ def main():
)
start_time = perf_counter()
- features, metadata, times, failures = track_method(
+ features, metadata, times = track_method(
inference_job,
writer,
live_tracking=args.live_tracking,
@@ -219,7 +230,6 @@ def main():
)
elapsed_time = end_time - start_time
write_tritonserver_metrics(args.metrics_endpoint, writer)
- assert len(failures) == 0, "should not be any failures"
number_of_tiles = len(features)
writer.add_scalar("number_of_tiles", number_of_tiles)
number_of_batches = math.ceil(number_of_tiles / args.batch_size)
diff --git a/benchmarking/triton_multiuser_benchmark.sh b/benchmarking/triton_multiuser_benchmark.sh
index 26c898b7..e86ff007 100755
--- a/benchmarking/triton_multiuser_benchmark.sh
+++ b/benchmarking/triton_multiuser_benchmark.sh
@@ -23,7 +23,7 @@ clear_cache() {
fi
}
-cmd="python ./benchmarking/triton_benchmark.py --numpy --inference-only --output ${output_dir} --model-name ${modelname} --wsi-path ${wsi_path} --batch-size 256 --url localhost:7985 --metrics-endpoint localhost:7986/metrics"
+cmd="python ./benchmarking/triton_benchmark.py --numpy --inference-only --output ${output_dir} --model-name ${modelname} --wsi-path ${wsi_path} --batch-size 128 --url localhost:7985 --metrics-endpoint localhost:7986/metrics"
if [[ "${modelname}" == "resnet50" || "${modelname}" == "resnet50_trt_uint8" || "${modelname}" == "gigapath_trt_uint8" ]]; then
cmd="${cmd} --nchw"
diff --git a/benchmarking/util.py b/benchmarking/util.py
index 42f37c4b..3e30ff7a 100644
--- a/benchmarking/util.py
+++ b/benchmarking/util.py
@@ -329,7 +329,7 @@ def parse_args():
if not args.wsi_path:
raise FileNotFoundError(f"No WSI files found in {args.wsi_path}")
- if not os.path.exists(args.output_path):
- raise FileNotFoundError(f"Did not find output directory {args.output_path}")
+ if not os.path.exists(args.output):
+ raise FileNotFoundError(f"Did not find output directory {args.output}")
return args
diff --git a/client.Dockerfile b/client.Dockerfile
index 16495094..e563cc7f 100644
--- a/client.Dockerfile
+++ b/client.Dockerfile
@@ -13,8 +13,8 @@ RUN echo 'APT::Install-Suggests "0";' >> /etc/apt/apt.conf.d/00-docker && \
rm -rf /var/lib/apt/lists/*
# install simple-triton
-WORKDIR /home/$USERNAME/code/simple_triton
-COPY simple_triton/ simple_triton
+WORKDIR /home/$USERNAME/code/triteia
+COPY triteia/ triteia
COPY pyproject.toml .
# comment out scm (i.e. git) line in pyproject.toml
RUN sed -i 's/.*\[tool.setuptools_scm\]/#&/g' pyproject.toml
@@ -42,8 +42,8 @@ RUN echo 'APT::Install-Suggests "0";' >> /etc/apt/apt.conf.d/00-docker && \
rm -rf /var/lib/apt/lists/*
USER $USERNAME
-WORKDIR /home/$USERNAME/simple_triton
-COPY --chown=$USERNAME:$USERNAME simple_triton/ simple_triton
+WORKDIR /home/$USERNAME/triteia
+COPY --chown=$USERNAME:$USERNAME triteia/ triteia
COPY --chown=$USERNAME:$USERNAME pyproject.toml .
@@ -73,8 +73,8 @@ RUN echo 'APT::Install-Suggests "0";' >> /etc/apt/apt.conf.d/00-docker && \
rm -rf /var/lib/apt/lists/*
RUN curl -fsSL https://get.docker.com | sh
# for jupyter notebooks as non-root
-RUN mkdir --mode a+rxw /.local /.jupyter /.cache /models/ /.config
-RUN chown $USERNAME:$USERNAME /home/$USERNAME/simple_triton/
+RUN mkdir --mode a+rxw /.local /.jupyter /.cache /.config
+RUN chown $USERNAME:$USERNAME /home/$USERNAME/triteia/
USER $USERNAME
COPY --chown=$USERNAME:$USERNAME README.md pyproject.toml ./
diff --git a/doc/overview_figure.pdf b/doc/overview_figure.pdf
new file mode 100644
index 00000000..d9422c0f
Binary files /dev/null and b/doc/overview_figure.pdf differ
diff --git a/doc/overview_figure.png b/doc/overview_figure.png
new file mode 100644
index 00000000..34356c22
Binary files /dev/null and b/doc/overview_figure.png differ
diff --git a/doc/overview_figure.tex b/doc/overview_figure.tex
new file mode 100644
index 00000000..92ed9108
--- /dev/null
+++ b/doc/overview_figure.tex
@@ -0,0 +1,123 @@
+\documentclass[tikz,border=10pt]{standalone}
+\usepackage{tikz}
+\usetikzlibrary{arrows.meta,positioning,fit,shadows.blur,calc}
+\usetikzlibrary{shapes.multipart}
+
+% --- Font: Linux Libertine (clean serif, excellent for scientific figures) ---
+\usepackage{libertine}
+\usepackage[libertine]{newtxmath}
+\usepackage[scaled]{helvet}
+\renewcommand\familydefault{\sfdefault}
+\usepackage[T1]{fontenc}
+
+\newcommand{\nestboxes}[4][]{
+ \begin{tikzpicture}[#1]
+ \foreach \L/\C[count=\n from 0] in {#4}{\fill[\C] (\n/2,#3-\n/2)node[black, below right]{\L} rectangle (#2-\n/2,\n/2);}
+ \end{tikzpicture}
+}
+
+\def\bW{0.22}
+\def\bH{0.07}
+%\draw[fill=blue!25, draw=blue!40, sharp corners] (\i*0.22,\i*0.07) rectangle ++(0.3,1.6+\i*0.07);
+
+\begin{document}
+\begin{tikzpicture}[
+ remember picture,
+ node distance=2.0cm and 2.2cm,
+ box/.style={draw, rounded corners, very thick, align=center, inner sep=12pt, fill=gray!3},
+ light/.style={fill=gray!6, draw=gray!60, rounded corners, thick},
+ title/.style={font=\bfseries\Huge},
+ arrow/.style={-{Latex[length=4mm]}, very thick},
+ thinarrow/.style={-{Latex[length=3mm]}, thick},
+ callout/.style={font=\Large\bfseries, align=center, text=gray!30!black},
+ every node/.style={font=\Large, text=black!100},
+]
+
+
+% --- Blocks ----------------------------------------------------
+\node[box, minimum width=30mm, minimum height=4mm] (data) {%
+ Whole-Slide Image(s)\\[-2mm]
+ \begin{tikzpicture}
+ \foreach \x in {1,...,3}
+ {
+ \node[yslant=-0.5] (wsi\x) at (\x*.55,0) {
+ \includegraphics[width=1.3cm, height=1.3cm]{wsi\x}
+ };
+ }
+ \end{tikzpicture}
+ \\[-15mm]
+};
+
+\node[box, below=3.0cm of data, minimum width=25mm, minimum height=28mm, fill=green!10] (client) {%
+ Python/Docker\\
+ environment\\[3mm]
+ \begin{tikzpicture}
+ \node[box, minimum width=55mm] (test) {\textbf{Triteia (client)} \\[4mm]
+ \begin{tikzpicture}
+ \node[draw=gray!60, rounded corners, thick, inner sep=8pt, minimum height=0pt] (aggreg) {Aggregation};
+ \end{tikzpicture}
+ };
+ \end{tikzpicture}
+ };
+
+\node[box, right=7.0cm of client, minimum height=38mm, fill=blue!10, inner sep=8pt, yshift=30mm] (server) {%
+ Server Hardware\\[-3mm]
+
+ \begin{tikzpicture}
+ \node[anchor=north, outer sep=0pt, inner sep=0pt, minimum width=0mm] (gpu-section) {
+ \hspace{-2mm}\Large GPU\\[4mm]
+ \begin{tikzpicture}[scale=1.0]
+ \draw[fill=gray!20, draw=gray!80, rounded corners=1pt] (0,0) rectangle (1.4,0.9);
+ \draw[fill=gray!70, draw=none] (0.1,0.2) rectangle (0.4,0.7);
+ \draw[fill=gray!60, draw=none] (1.1,0.2) rectangle (1.3,0.7);
+ % connector lines
+ \draw[gray!60, thick] (1.4,0.45) -- ++(0.25,0);
+ \end{tikzpicture}
+ };
+
+ \node[outer sep=0pt, inner sep=0pt, anchor=north] (cpu-section) at (3.6, 0.0) {
+ \\[-1.5mm]
+ \Large CPU\\[2mm]
+ \begin{tikzpicture}[]
+ \draw[fill=gray!15, draw=gray!80, thick] (0.1,0) rectangle (1.0,0.9);
+ \foreach \x in {0.1,0.3,0.5,0.7,0.9}{
+ \draw[gray!100] (\x+0.1,-0.1) -- ++(0,-0.1);
+ \draw[gray!100] (\x+0.1,1.0) -- ++(0,0.1);
+ \draw[gray!100] (0.0,\x) -- ++(-0.1,0);
+ \draw[gray!100] (1.1,\x) -- ++(0.1,0);
+ }
+ \end{tikzpicture}
+ };
+ \end{tikzpicture}\\[-5mm]
+
+ \begin{tikzpicture}
+
+ \node[box, fill=green!10] (test2) {Docker environment\\[3mm]
+ \begin{tikzpicture}
+ \node[box] (test3) {\textbf{Triteia (server)} \\[5mm]
+ \begin{tikzpicture}[scale=0.1, baseline={(current bounding box.center)}]
+ \node[box, minimum height=1pt] (nvidiatri) {NVIDIA Triton\\inference server
+ };
+ \end{tikzpicture}\\[1mm]
+
+ \begin{tikzpicture}[scale=0.8, baseline={(current bounding box.center)}]
+ \node[inner sep=0] (models) at (2,0.8) { hosts DL model(s)\\[4mm]
+ \begin{tikzpicture}
+ \foreach \i in {0,1,2,3} {
+ \draw[fill=blue!25, draw=blue!40, sharp corners] (\i*\bW,\i*\bH) rectangle ++(\bW*1.3,\bH*15+\i*\bH);
+ }
+ \end{tikzpicture}
+ };
+ \end{tikzpicture}\\[-3mm]
+ };
+ \end{tikzpicture}
+ };
+ \end{tikzpicture}
+};
+
+% --- Flow arrows ------------------------------------------------
+\draw[arrow] (data) -- node[left, yshift=2mm, callout] {image tiles} (client);
+\draw[arrow] (client.east |- nvidiatri.west) -- node[midway, above=0.2cm, callout] {requests (gRPC)} (nvidiatri.west);
+\draw[arrow] (server.west |- aggreg.east) -- node[midway, above right=0.2cm and -10.0mm, callout] {results (gRPC)} (aggreg.east);
+\end{tikzpicture}
+\end{document}
diff --git a/doc/wsi1.png b/doc/wsi1.png
new file mode 100644
index 00000000..b23bbe32
Binary files /dev/null and b/doc/wsi1.png differ
diff --git a/doc/wsi2.png b/doc/wsi2.png
new file mode 100644
index 00000000..39c49ecd
Binary files /dev/null and b/doc/wsi2.png differ
diff --git a/doc/wsi3.png b/doc/wsi3.png
new file mode 100644
index 00000000..c7bc2a03
Binary files /dev/null and b/doc/wsi3.png differ
diff --git a/examples/export_pytorch_model.ipynb b/examples/export_pytorch_model.ipynb
index e8c0cf3c..68a73a90 100644
--- a/examples/export_pytorch_model.ipynb
+++ b/examples/export_pytorch_model.ipynb
@@ -5,9 +5,9 @@
"id": "2c6e72468de23815",
"metadata": {},
"source": [
- "# Exporting a PyTorch model for use with simple_triton\n",
+ "# Exporting a PyTorch model for use with triteia\n",
"\n",
- "In this notebook, we will export a PyTorch model for use with simple_triton/tritonserver with a PyTorch backend.\n",
+ "In this notebook, we will export a PyTorch model for use with triteia/tritonserver with a PyTorch backend.\n",
"\n",
"We will:\n",
"1. Create an example PyTorch model (or provide your own model weights)\n",
@@ -50,9 +50,9 @@
"import torch\n",
"from jupyter_core.utils import ensure_dir_exists\n",
"\n",
- "from simple_triton.feature_extraction import inference\n",
- "from simple_triton.utils import analyze\n",
- "from simple_triton.model import TritonModel\n"
+ "from triteia.feature_extraction import inference\n",
+ "from triteia.utils import analyze\n",
+ "from triteia.model import TritonModel\n"
],
"outputs": [
{
diff --git a/examples/feature_extraction.ipynb b/examples/feature_extraction.ipynb
index fad0c4d1..e34d6859 100644
--- a/examples/feature_extraction.ipynb
+++ b/examples/feature_extraction.ipynb
@@ -104,7 +104,7 @@
}
},
"source": [
- "from simple_triton.encoders import tf_encoder\n",
+ "from triteia.encoders import tf_encoder\n",
"\n",
"# model parameters\n",
"keras_name = \"EfficientNetV2S\"\n",
@@ -849,7 +849,7 @@
},
"source": [
"from pprint import pprint\n",
- "from simple_triton.config import TensorflowConfig\n",
+ "from triteia.config import TensorflowConfig\n",
"\n",
"# build a basic configuration specifying only maximum batch size and model name\n",
"max_batch_size = 64\n",
@@ -889,7 +889,7 @@
}
},
"source": [
- "from simple_triton.model import TritonModel\n",
+ "from triteia.model import TritonModel\n",
"\n",
"# load tensorflow model - set maximum batch size\n",
"model = TritonModel(model_name, \"localhost:8001\")\n",
@@ -944,7 +944,7 @@
}
},
"source": [
- "from simple_triton.config import (\n",
+ "from triteia.config import (\n",
" InstanceGroup,\n",
" TensorflowMixedPrecision,\n",
" TensorflowOptimization,\n",
@@ -1005,9 +1005,9 @@
},
"source": [
"import numpy as np\n",
- "from simple_triton.feature_extraction import inference, study\n",
- "from simple_triton.tile_iterators import TiffPrefetch\n",
- "from simple_triton.utils import analyze\n",
+ "from triteia.feature_extraction import inference, study\n",
+ "from triteia.tile_iterators import TiffPrefetch\n",
+ "from triteia.utils import analyze\n",
"from time import time\n",
"\n",
"# slide parameters\n",
@@ -1525,8 +1525,8 @@
}
],
"source": [
- "from simple_triton.io.tfr_reader import read_record, peek\n",
- "from simple_triton.io.tfr_writer import write_record\n",
+ "from triteia.io.tfr_reader import read_record, peek\n",
+ "from triteia.io.tfr_writer import write_record\n",
"\n",
"# create dummy labels\n",
"labels = {\"labels\": np.random.uniform(size=(10))}\n",
diff --git a/examples/patch_inference.ipynb b/examples/patch_inference.ipynb
index 9c753d72..0fa16f8e 100644
--- a/examples/patch_inference.ipynb
+++ b/examples/patch_inference.ipynb
@@ -53,9 +53,9 @@
"import matplotlib.pyplot as plt\n",
"from pprint import pprint\n",
"\n",
- "from simple_triton.feature_extraction import inference\n",
- "from simple_triton.utils import analyze\n",
- "from simple_triton.model import TritonModel"
+ "from triteia.feature_extraction import inference\n",
+ "from triteia.utils import analyze\n",
+ "from triteia.model import TritonModel"
],
"outputs": [
{
@@ -175,27 +175,23 @@
"execution_count": 3
},
{
- "cell_type": "markdown",
- "id": "ca78525d723d5396",
"metadata": {},
+ "cell_type": "markdown",
"source": [
- "## Create a dataloader compatible with simple_triton\n",
+ "## Create a dataloader compatible with triteia\n",
"In this code, we use MONAI, which is a wrapper around PyTorch. I like to use MONAI since:\n",
" * they have many common image transform functions\n",
" * Their dataloader is easy to use.\n",
"\n",
"You do not have use MONAI. The only important thing is to have an iterator that returns (sample, metadata) pairs."
- ]
+ ],
+ "id": "6cf52a1d819f225b"
},
{
+ "metadata": {},
"cell_type": "code",
- "id": "ebc89a3e6f1cbc9c",
- "metadata": {
- "ExecuteTime": {
- "end_time": "2026-03-04T12:36:08.669024666Z",
- "start_time": "2026-03-04T12:36:08.582815308Z"
- }
- },
+ "outputs": [],
+ "execution_count": null,
"source": [
"transformations = mt.Compose(\n",
" [\n",
@@ -206,7 +202,7 @@
"\n",
"def triton_collate_fn(batch):\n",
" '''\n",
- " collate function for simple_triton inference.\n",
+ " collate function for triteia inference.\n",
" This function changes the MONAI dataloader to return (sample, metadata)\n",
" instead of a dictionary, which is the default behaviour.\n",
" '''\n",
@@ -216,8 +212,7 @@
"\n",
"dl = DataLoader(dataset=Dataset(all_data, transformations), batch_size=32, shuffle=False, collate_fn=triton_collate_fn)\n"
],
- "outputs": [],
- "execution_count": 4
+ "id": "e77d29020b862204"
},
{
"cell_type": "markdown",
diff --git a/examples/resnet50_trt/export_model.sh b/examples/resnet50_trt/export_model.sh
index ddbea30c..14786c64 100755
--- a/examples/resnet50_trt/export_model.sh
+++ b/examples/resnet50_trt/export_model.sh
@@ -16,7 +16,7 @@ docker run -it --gpus all -w /trt_optimize -v $PWD:/trt_optimize nvcr.io/nvidia/
set +xe
model_dst_dir="${git_root_dir}/models/resnet50_trt_uint8"
-test -d "${model_dst_dir}" || mkdir -p "${model_dst_dir}"
+test -d "${model_dst_dir}/1" || mkdir -p "${model_dst_dir}/1"
dst="${model_dst_dir}/config.pbtxt"
cp -v config.pbtxt "${dst}"
# we need to replace TYPE_FP32 with UINT8. But not for the model outputs.
@@ -34,7 +34,7 @@ docker run -it --gpus all -w /resnet50_eg -v $PWD:/resnet50_eg nvcr.io/nvidia/py
docker run -it --gpus all -w /trt_optimize -v $PWD:/trt_optimize nvcr.io/nvidia/tensorrt:${nvidia_version} trtexec --onnx=resnet50_fp32.onnx --saveEngine=modelfp32.plan --useCudaGraph --minShapes=input_0:1x3x224x224 --optShapes=input_0:16x3x224x224 --maxShapes=input_0:256x3x224x224
set +xe
model_dst_dir="${git_root_dir}/models/resnet50_trt_float32"
-test -d "${model_dst_dir}" || mkdir -p "${model_dst_dir}"
+test -d "${model_dst_dir}/1" || mkdir -p "${model_dst_dir}/1"
dst="${model_dst_dir}/config.pbtxt"
cp config.pbtxt "${dst}"
cp -v modelfp32.plan "${model_dst_dir}/1/model.plan"
diff --git a/launch_test_container.sh b/launch_test_container.sh
index 9a6806a1..7d8b91c2 100755
--- a/launch_test_container.sh
+++ b/launch_test_container.sh
@@ -11,7 +11,7 @@ if [ -z "${HF_TOKEN}" ]; then
printf "You can add it by quitting this script/container, and then \`export HF_TOKEN=...\` and re-start the script\n"
fi
HF_TOKEN=${HF_TOKEN:-undefined}
-model_tmp_directory=/tmp/simple_triton_test_data_$USER
+model_tmp_directory=/tmp/triteia_test_data_$USER
test -d "${model_tmp_directory}" || (mkdir "${model_tmp_directory}" && chmod -R 777 "${model_tmp_directory}")
cp -r models/* "${model_tmp_directory}"/ || exit 1
-docker run --security-opt seccomp:unconfined --network=host -v "${model_tmp_directory}:${model_tmp_directory}:rw" -e TRITON_TMP_DIR="${model_tmp_directory}/" -e HF_TOKEN=${HF_TOKEN} --shm-size=2g -v /var/run/docker.sock:/var/run/docker.sock --rm --name tritonclient_test -it simple_triton_client:latest
+docker run --security-opt seccomp:unconfined --network=host -v "${model_tmp_directory}:${model_tmp_directory}:rw" -e TRITON_TMP_DIR="${model_tmp_directory}/" -e HF_TOKEN=${HF_TOKEN} --shm-size=2g -v /var/run/docker.sock:/var/run/docker.sock --rm --name triteia_test -it triteia:latest
diff --git a/models/gigapath/model.py b/models/gigapath/model.py
index 970e72e8..34f63ebe 100644
--- a/models/gigapath/model.py
+++ b/models/gigapath/model.py
@@ -6,8 +6,19 @@
import torch
import triton_python_backend_utils as pb_utils
import tritonclient.utils as triton_utils
-from huggingface_hub import login
+from huggingface_hub import login, get_token
+
+def require_hf_token(repo_id):
+ token = get_token()
+ if not token:
+ raise RuntimeError(
+ "Missing Hugging Face access token for "
+ f"{repo_id}. Set HF_TOKEN to a token with read access to the repository."
+ )
+ os.environ["HF_TOKEN"] = token
+ os.environ.setdefault("HUGGING_FACE_HUB_TOKEN", token)
+ return token
class TritonPythonModel:
@staticmethod
@@ -57,9 +68,7 @@ def initialize(self, args):
self.output0_dtype = pb_utils.triton_string_to_numpy(
output0_config["data_type"]
)
- login(
- os.getenv("HF_TOKEN")
- ) # User Access Token, found at https://huggingface.co/settings/tokens
+ require_hf_token("hf_hub:prov-gigapath/prov-gigapath")
self.model = timm.create_model(
"hf_hub:prov-gigapath/prov-gigapath",
pretrained=True,
diff --git a/models/resnet50/model.py b/models/resnet50/model.py
new file mode 100644
index 00000000..52479294
--- /dev/null
+++ b/models/resnet50/model.py
@@ -0,0 +1,119 @@
+# Link: https://github.com/triton-inference-server/python_backend/tree/main/examples/instance_kind
+
+# Copyright 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions
+# are met:
+# * Redistributions of source code must retain the above copyright
+# notice, this list of conditions and the following disclaimer.
+# * Redistributions in binary form must reproduce the above copyright
+# notice, this list of conditions and the following disclaimer in the
+# documentation and/or other materials provided with the distribution.
+# * Neither the name of NVIDIA CORPORATION nor the names of its
+# contributors may be used to endorse or promote products derived
+# from this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
+# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
+# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
+# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
+# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+import json
+
+import torch
+import triton_python_backend_utils as pb_utils
+from torch.utils.dlpack import to_dlpack
+from transformers import AutoImageProcessor, ResNetModel, AutoModel
+
+
+class TritonPythonModel:
+ def initialize(self, args):
+ """
+ This function initializes pre-trained ResNet50 model,
+ depending on the value specified by an `instance_group` parameter
+ in `config.pbtxt`.
+
+ Depending on what `instance_group` was specified in
+ the config.pbtxt file (KIND_CPU or KIND_GPU), the model instance
+ will be initialised on a cpu, a gpu, or both. If `instance_group` was
+ not specified in the config file, then models will be loaded onto
+ the default device of the framework.
+ """
+ # Here we set up the device onto which our model will beloaded,
+ # based on specified `model_instance_kind` and `model_instance_device_id`
+ # fields.
+ device = "cuda" if args["model_instance_kind"] == "GPU" else "cpu"
+ device_id = args["model_instance_device_id"]
+ self.device = f"{device}:{device_id}"
+ self.processor = AutoImageProcessor.from_pretrained("microsoft/resnet-50", use_fast=True)
+
+ self.model = ResNetModel.from_pretrained("microsoft/resnet-50") \
+ .to(self.device) \
+ .eval()
+
+ self.model_config = model_config = json.loads(args["model_config"])
+ output0_config = pb_utils.get_output_config_by_name(model_config, "output_0")
+ self.gpu_id = args.get("model_instance_device_id", 0)
+ self.output0_dtype = pb_utils.triton_string_to_numpy(
+ output0_config["data_type"]
+ )
+
+ @staticmethod
+ def auto_complete_config(model_config):
+ """Returns a minimal model configuration for the uni model.
+
+ Parameters
+ ----------
+ model_config : pb_utils.ModelConfig
+ An object containing the existing model configuration.
+
+ Returns
+ -------
+ pb_utils.ModelConfig
+ An object containing the auto-completed model configuration
+ """
+ inputs = [
+ {
+ "name": "input_0",
+ "data_type": "TYPE_UINT8",
+ "dims": [3, 224, 224],
+ }
+ ]
+ outputs = [{"name": "output_0", "data_type": "TYPE_FP32", "dims": [1000]}]
+ config = model_config.as_dict()
+ input_names = [i["name"] for i in config["input"]]
+ output_names = [i["name"] for i in config["output"]]
+ for i in inputs:
+ if i["name"] not in input_names:
+ model_config.add_input(i)
+ for o in outputs:
+ if o["name"] not in output_names:
+ model_config.add_output(o)
+ model_config.set_max_batch_size(256)
+
+ return model_config
+
+ def execute(self, requests):
+ """
+ This function receives a list of requests (`pb_utils.InferenceRequest`),
+ performs inference on every request and appends it to responses.
+ """
+ responses = [None] * len(requests)
+
+ with torch.inference_mode():
+ for i, request in enumerate(requests):
+ inputs = self.processor(pb_utils.get_input_tensor_by_name(request, "input_0").as_numpy(),
+ return_tensors="pt",
+ device=self.device)
+ result = self.model(**inputs).pooler_output.flatten(1, -1)
+ out_tensor = pb_utils.Tensor.from_dlpack("output_0", to_dlpack(result))
+ responses[i] = pb_utils.InferenceResponse([out_tensor])
+ return responses
diff --git a/models/uni/model.py b/models/uni/model.py
index 2e549630..2337fc3d 100644
--- a/models/uni/model.py
+++ b/models/uni/model.py
@@ -1,12 +1,25 @@
import json
-import os
import numpy as np
import timm
+import os
import torch
import triton_python_backend_utils as pb_utils
import tritonclient.utils as triton_utils
-from huggingface_hub import login
+from huggingface_hub import get_token
+
+def require_hf_token(repo_id):
+
+ token = get_token()
+ if not token:
+ raise RuntimeError(
+ "Missing Hugging Face access token for "
+ f"{repo_id}. Set HF_TOKEN to a token with read access to the repository."
+ )
+
+ os.environ["HF_TOKEN"] = token
+ os.environ.setdefault("HUGGING_FACE_HUB_TOKEN", token)
+ return token
class TritonPythonModel:
@@ -57,9 +70,7 @@ def initialize(self, args):
self.output0_dtype = pb_utils.triton_string_to_numpy(
output0_config["data_type"]
)
- login(
- os.getenv("HF_TOKEN")
- ) # User Access Token, found at https://huggingface.co/settings/tokens
+ require_hf_token("MahmoodLab/UNI")
self.model = timm.create_model(
"hf-hub:MahmoodLab/uni",
pretrained=True,
diff --git a/pyproject.toml b/pyproject.toml
index 4f1b1d29..efac89f7 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -3,7 +3,7 @@ requires = ["setuptools", "setuptools-scm"]
build-backend = "setuptools.build_meta"
[project]
-name = "simple_triton"
+name = "triteia"
authors = [
{name = "Lee Cooper", email = "lee.cooper@northwestern.edu"},
]
@@ -36,6 +36,7 @@ examples = [
"matplotlib",
"medmnist@git+https://github.com/andsild/MedMNIST-no-gpu",
"monai",
+ "onnxscript",
"pooch",
"scikit-learn",
"umap-learn",
@@ -44,14 +45,14 @@ examples = [
]
[project.scripts]
-inference = "simple_triton.feature_extraction:main"
+inference = "triteia.feature_extraction:main"
[tool.setuptools]
-packages = ["simple_triton", "simple_triton.io"]
+packages = ["triteia", "triteia.io"]
[tool.setuptools_scm]
[project.urls]
-"Github" = "https://github.com/PathologyDataScience/simple_triton"
+"Github" = "https://github.com/PathologyDataScience/triteia"
diff --git a/tests/test_config.py b/tests/test_config.py
index c6292945..06422a49 100644
--- a/tests/test_config.py
+++ b/tests/test_config.py
@@ -7,7 +7,7 @@
from .triton import triton
from .data import data
-from simple_triton.config import (
+from triteia.config import (
InstanceGroup,
ModelInput,
ModelOutput,
@@ -17,7 +17,7 @@
TensorflowXla,
TensorRt,
)
-from simple_triton.model import TritonModel, create_client
+from triteia.model import TritonModel, create_client
MODEL = "EfficientNetV2S.tensorflow"
BASIC = {"name": MODEL}
diff --git a/tests/test_inference.py b/tests/test_inference.py
index a97f6053..301b72fb 100644
--- a/tests/test_inference.py
+++ b/tests/test_inference.py
@@ -1,10 +1,10 @@
from .data import data, hash_inference, inferred, it_kwargs_icc
import numpy as np
import os
-from simple_triton.config import PythonConfig, TensorflowConfig
-from simple_triton.feature_extraction import inference
-from simple_triton.model import TritonModel
-from simple_triton.tile_iterators import TiffPrefetch
+from triteia.config import PythonConfig, TensorflowConfig
+from triteia.feature_extraction import inference
+from triteia.model import TritonModel
+from triteia.tile_iterators import TiffPrefetch
from .triton import triton
diff --git a/tests/test_model.py b/tests/test_model.py
index 7700e42b..768fb2fd 100644
--- a/tests/test_model.py
+++ b/tests/test_model.py
@@ -3,7 +3,7 @@
from multiprocessing import get_context
import numpy as np
import pytest
-from simple_triton.model import TritonModel, create_client
+from triteia.model import TritonModel, create_client
import sys
import time
import tritonclient.grpc as grpcclient
diff --git a/tests/test_tile_iterators.py b/tests/test_tile_iterators.py
index 35479e8f..40f58c02 100644
--- a/tests/test_tile_iterators.py
+++ b/tests/test_tile_iterators.py
@@ -7,7 +7,7 @@
tiles_icc,
)
import numpy as np
-from simple_triton.tile_iterators import TiffPrefetch
+from triteia.tile_iterators import TiffPrefetch
def compare_dict(x, y):
diff --git a/tox.ini b/tox.ini
index b110c21d..86deef37 100644
--- a/tox.ini
+++ b/tox.ini
@@ -39,15 +39,15 @@ testpaths =
[coverage:run]
concurrency = multiprocessing
include =
- simple_triton/*
- {envsitepackagesdir}/simple_triton/*
+ triteia/*
+ {envsitepackagesdir}/triteia/*
omit =
tests/*
[coverage:paths]
source =
- simple_triton/
- {envsitepackagesdir}/simple_triton/
+ triteia/
+ {envsitepackagesdir}/triteia/
[coverage:html]
directory = .tox/coverage/coverage.html
diff --git a/simple_triton/__init__.py b/triteia/__init__.py
similarity index 100%
rename from simple_triton/__init__.py
rename to triteia/__init__.py
diff --git a/simple_triton/config.py b/triteia/config.py
similarity index 93%
rename from simple_triton/config.py
rename to triteia/config.py
index 55ccd27e..6b8ec39f 100644
--- a/simple_triton/config.py
+++ b/triteia/config.py
@@ -430,6 +430,60 @@ def __init__(
self.config.update(xla.config)
+class TensorRTConfig(PythonConfig):
+ """A model configuration for the tensorflow backend.
+
+ This class can generate JSON format dictionaries for use with model loading
+ functions, and can save configurations in protocol buffer format for
+ file-based configuration.
+
+ Parameters
+ ----------
+ name : str
+ Model name as stored in the model repository.
+ input : ModelInput or list
+ Model inputs.
+ output : ModelOutput or list
+ Model outputs.
+ instance_group : InstanceGroup
+ An instance group configuration defining model resources.
+ max_batch_size : int
+ The maximum number of samples in a request. Use 0 for a non-batching model.
+ optimization : TensorflowOptimization
+ Python backend optimization configuration. Default value None enables
+ pinned memory by default.
+ response_cache : bool
+ Whether to cache model input-output pairs. See reference below. Default value
+ is False for no caching.
+
+ References
+ ----------
+ https://github.com/triton-inference-server/server/blob/main/docs/user_guide/response_cache.md
+ """
+
+ def __init__(
+ self,
+ name,
+ max_batch_size,
+ input=None,
+ output=None,
+ instance_group=None,
+ dynamic_batching=None,
+ response_cache=False,
+ ):
+ super(TensorRTConfig, self).__init__(
+ name=name,
+ input=input,
+ output=output,
+ instance_group=instance_group,
+ max_batch_size=max_batch_size,
+ dynamic_batching=dynamic_batching,
+ response_cache=response_cache,
+ )
+ self.config["backend"] = "tensorrt"
+ self.config["platform"] = "tensorrt_plan"
+
+
class TensorflowConfig(PythonConfig):
"""A model configuration for the tensorflow backend.
diff --git a/simple_triton/encoders.py b/triteia/encoders.py
similarity index 100%
rename from simple_triton/encoders.py
rename to triteia/encoders.py
diff --git a/simple_triton/feature_extraction.py b/triteia/feature_extraction.py
similarity index 99%
rename from simple_triton/feature_extraction.py
rename to triteia/feature_extraction.py
index bd388fc5..b3ef4635 100644
--- a/simple_triton/feature_extraction.py
+++ b/triteia/feature_extraction.py
@@ -23,10 +23,10 @@
from pytriton.client import ModelClient
from tqdm import tqdm
-from simple_triton.io.tfr_writer import write_record
-from simple_triton.model import TritonModel
-from simple_triton.tile_iterators import TiffPrefetch
-from simple_triton.utils import (
+from triteia.io.tfr_writer import write_record
+from triteia.model import TritonModel
+from triteia.tile_iterators import TiffPrefetch
+from triteia.utils import (
analyze,
init_tb_writer,
track_method,
diff --git a/simple_triton/io/__init__.py b/triteia/io/__init__.py
similarity index 100%
rename from simple_triton/io/__init__.py
rename to triteia/io/__init__.py
diff --git a/simple_triton/io/tfr_reader.py b/triteia/io/tfr_reader.py
similarity index 98%
rename from simple_triton/io/tfr_reader.py
rename to triteia/io/tfr_reader.py
index 3c670d72..368a619f 100644
--- a/simple_triton/io/tfr_reader.py
+++ b/triteia/io/tfr_reader.py
@@ -2,8 +2,8 @@
import tensorflow as tf
-from simple_triton.io import slide_keys, tile_keys
-from simple_triton.io.tfr_transforms import flatten, structure
+from triteia.io import slide_keys, tile_keys
+from triteia.io.tfr_transforms import flatten, structure
def peek(serialized):
diff --git a/simple_triton/io/tfr_transforms.py b/triteia/io/tfr_transforms.py
similarity index 100%
rename from simple_triton/io/tfr_transforms.py
rename to triteia/io/tfr_transforms.py
diff --git a/simple_triton/io/tfr_writer.py b/triteia/io/tfr_writer.py
similarity index 99%
rename from simple_triton/io/tfr_writer.py
rename to triteia/io/tfr_writer.py
index ad3045b0..47cd21a2 100644
--- a/simple_triton/io/tfr_writer.py
+++ b/triteia/io/tfr_writer.py
@@ -3,9 +3,9 @@
import numpy as np
import tensorflow as tf
-from simple_triton.io import slide_keys, tile_keys
-from simple_triton.io.tfr_reader import peek, read_record
-from simple_triton.io.tfr_transforms import structure
+from triteia.io import slide_keys, tile_keys
+from triteia.io.tfr_reader import peek, read_record
+from triteia.io.tfr_transforms import structure
# acceptable types for user-provided metadata
variable_type_list = [bytes, int, float, str, bool]
diff --git a/simple_triton/model.py b/triteia/model.py
similarity index 100%
rename from simple_triton/model.py
rename to triteia/model.py
diff --git a/simple_triton/tile_iterators.py b/triteia/tile_iterators.py
similarity index 100%
rename from simple_triton/tile_iterators.py
rename to triteia/tile_iterators.py
diff --git a/simple_triton/utils.py b/triteia/utils.py
similarity index 100%
rename from simple_triton/utils.py
rename to triteia/utils.py
diff --git a/triton_overview.png b/triton_overview.png
deleted file mode 100644
index 9c33ab40..00000000
Binary files a/triton_overview.png and /dev/null differ