Tracing Data Flow in Hyrax#
When working with a new dataset or model it can be hard to know whether your data is being processed correctly. Hyrax provides a trace mode that lets you follow a small batch of data items through the entire pipeline — from raw dataset access, through batching and input preparation, all the way to model evaluation — so you can inspect what is actually happening at each step.
This notebook walks through a typical trace session:
Setting up a Hyrax instance
Running a verb with
trace=Nto capture pipeline dataPrinting and navigating the returned
TraceResultDrilling into individual stages and function calls
Trace mode is intended for interactive use in notebooks — it is not a production profiling or logging tool.
Install Hyrax#
Skip this step if you have already installed Hyrax in your environment.
[1]:
# %pip install hyrax
Create a Hyrax instance#
We start by creating a Hyrax instance and pointing it at a simple built-in model and dataset. Here we use HyraxAutoencoder trained on HyraxRandomDataset — a dataset that generates random tensors without requiring any downloaded data, which makes it convenient for experimentation.
You can replace this with your own model and dataset configuration to trace your real workflow.
[2]:
import hyrax
h = hyrax.Hyrax()
Configure the model and data#
We will use the HyraxAutoencoder model and HyraxRandomDataset to keep this notebook self-contained. The random dataset produces small random tensors ([1, 32, 32]) that resemble single-band image cutouts.
[3]:
h.config["model"]["name"] = "HyraxAutoencoder"
h.config["model"]["HyraxAutoencoder"]["input_shape"] = [1, 32, 32]
# Use a small 1-channel 32×32 image shape to keep things fast
h.config["data_set"]["HyraxRandomDataset"]["shape"] = [1, 32, 32]
h.config["data_set"]["HyraxRandomDataset"]["size"] = 50
h.config["data_set"]["HyraxRandomDataset"]["seed"] = 42
data_request = {
"train": {
"data": {
"dataset_class": "HyraxRandomDataset",
"data_location": "./trace_data",
"fields": ["image"],
"primary_id_field": "object_id",
},
},
"infer": {
"data": {
"dataset_class": "HyraxRandomDataset",
"data_location": "./trace_data",
"fields": ["image"],
"primary_id_field": "object_id",
},
},
}
h.set_config("data_request", data_request)
[4]:
from hyrax.datasets.random.hyrax_random_dataset import HyraxRandomDataset
import numpy as np
def collate_image(self, samples: list[dict]) -> dict:
collated_data = {}
collated_data["image"] = np.stack([sample["image"] for sample in samples], axis=0)
collated_data["mask"] = np.zeros_like(collated_data["image"], dtype=bool)
return collated_data
HyraxRandomDataset.collate = collate_image
We attach either a collate_image function to HyraxRandomDataset above. We could instead attach a dataset-level collate function like the one shown below. The TraceResult will contain information for whichever function is attached. Field-level collation functions are described in the field_level_collation group which will be empty if a dataset-level collate function is defined. Dataset-level collate functions are described in the collate group.
def collate(self, samples: list[dict]) -> dict:
collated_data = {}
collated_data["image"] = np.stack([sample["image"] for sample in samples], axis=0)
collated_data["mask"] = np.zeros_like(collated_data["image"], dtype=bool)
return collated_data
HyraxRandomDataset.collate = collate
Running a verb in trace mode#
Any instrumented verb (train, infer, test) accepts a trace=N keyword argument. N controls how many data items are traced through the pipeline — keep this small (2–10) to get a readable output.
When trace=N is passed, the verb:
Processes only a single batch of
Nitems instead of the full datasetReturns a
TraceResultobject instead of the usual verb return value
The TraceResult captures the inputs and outputs of every major pipeline step.
[5]:
trace_result = h.train(trace=2)
[2026-10-28 12:25:02,960 hyrax.trace:WARNING] Starting Trace
[2026-10-28 12:25:02,961 hyrax.trace:WARNING] Trace mode enabled, will only run a single batch of length 2
[2026-10-28 12:25:02,979 hyrax.models.model_registry:INFO] Setting model's self.optimizer from config: torch.optim.SGD with arguments: {'lr': 0.01, 'momentum': 0.9}.
[2026-10-28 12:25:02,979 hyrax.models.model_registry:INFO] Setting model's self.criterion from config: torch.nn.CrossEntropyLoss with default arguments.
[2026-10-28 12:25:02,980 hyrax.models.model_registry:INFO] Setting model's self.scheduler from config: torch.optim.lr_scheduler.ExponentialLR
with arguments: {'gamma': 1}.
[2026-10-28 12:25:03,045 hyrax.verbs.train:INFO] Using a single process for training.
[2026-10-28 12:25:03,045 hyrax.verbs.train:INFO] Training model: HyraxAutoencoder
[2026-10-28 12:25:03,046 hyrax.verbs.train:INFO] Training dataset(s):
{'train': Name: data (primary dataset)
Dataset class: HyraxRandomDataset
Data location: /home/drew/code/hyrax/docs/pre_executed/trace_data
Selected items: 50
Primary ID field: object_id
Requested fields: image
}
2026-10-28 12:25:03,091 ignite.distributed.auto.auto_dataloader INFO: Use data loader kwargs for dataset '<torch.utils.data.da':
{'sampler': <torch.utils.data.sampler.WeightedRandomSampler object at 0x723aa6d10c20>, 'batch_size': 2, 'collate_fn': <bound method CollationMixin.collate of Name: data (primary dataset)
Dataset class: HyraxRandomDataset
Data location: /home/drew/code/hyrax/docs/pre_executed/trace_data
Selected items: 50
Primary ID field: object_id
Requested fields: image
>, 'pin_memory': True}
2026/10/28 12:25:03 INFO mlflow.system_metrics.system_metrics_monitor: Skip logging GPU metrics. Set logger level to DEBUG for more details.
2026/10/28 12:25:03 INFO mlflow.system_metrics.system_metrics_monitor: Started monitoring system metrics.
[2026-10-28 12:25:04,856 hyrax.pytorch_ignite:INFO] Total training time: 1.02[s]
2026/10/28 12:25:04 INFO mlflow.system_metrics.system_metrics_monitor: Stopping system metrics monitoring...
2026/10/28 12:25:04 INFO mlflow.system_metrics.system_metrics_monitor: Successfully terminated system metrics monitoring!
[2026-10-28 12:25:04,882 hyrax.verbs.train:INFO] Finished Training
Printing the trace#
Printing trace_result gives a high-level summary of all pipeline stages and the function calls captured within them. Each entry shows the function name, its input and output names, the shapes / hashes of any tensors, and how long the call took.
[6]:
print(trace_result)
Trace Stages {
dataset_getter: [
data__get_image(index) -> image duration=0.00381 ms
inputs:
index = 6
outputs:
image = <numpy.ndarray shape=(1, 32, 32) hash=66186351732588544 device=cpu>
data__get_object_id(index) -> object_id duration=0.00378 ms
inputs:
index = 6
outputs:
object_id = '29'
data__get_image(index) -> image duration=0.00177 ms
inputs:
index = 16
outputs:
image = <numpy.ndarray shape=(1, 32, 32) hash=28339009725923328 device=cpu>
data__get_object_id(index) -> object_id duration=0.000736 ms
inputs:
index = 16
outputs:
object_id = '39'
]
resolve_data: [
DataProvider__resolve_data(index) -> data_dict duration=0.034 ms
inputs:
index = 6
outputs:
data_dict = {
data: {
image: <numpy.ndarray shape=(1, 32, 32) hash=66186351732588544 device=cpu>
}
object_id: '29'
}
DataProvider__resolve_data(index) -> data_dict duration=0.01 ms
inputs:
index = 16
outputs:
data_dict = {
data: {
image: <numpy.ndarray shape=(1, 32, 32) hash=28339009725923328 device=cpu>
}
object_id: '39'
}
]
field_level_collation: []
collate: [
DataProvider__collate(batch_dicts) -> batch_dict duration=0.124 ms
inputs:
batch_dicts = <list len=2> [
{
data: {
image: <numpy.ndarray shape=(1, 32, 32) hash=66186351732588544 device=cpu>
}
object_id: '29'
}
{
data: {
image: <numpy.ndarray shape=(1, 32, 32) hash=28339009725923328 device=cpu>
}
object_id: '39'
}
]
outputs:
batch_dict = {
object_id: <numpy.ndarray shape=(2,) hash=14736351079544187634 device=cpu>
data: {
image: <numpy.ndarray shape=(2, 1, 32, 32) hash=40402718698766336 device=cpu>
mask: <numpy.ndarray shape=(2, 1, 32, 32) hash=0 device=cpu>
}
}
data__collate(samples) -> batch_dict duration=0.0376 ms
inputs:
samples = <list len=2> [
{
image: <numpy.ndarray shape=(1, 32, 32) hash=66186351732588544 device=cpu>
}
{
image: <numpy.ndarray shape=(1, 32, 32) hash=28339009725923328 device=cpu>
}
]
outputs:
batch_dict = {
image: <numpy.ndarray shape=(2, 1, 32, 32) hash=40402718698766336 device=cpu>
mask: <numpy.ndarray shape=(2, 1, 32, 32) hash=0 device=cpu>
}
DataProvider__handle_nans(batch_dict) -> batch_dict_no_nan duration=0.0665 ms
inputs:
batch_dict = {
object_id: <numpy.ndarray shape=(2,) hash=14736351079544187634 device=cpu>
data: {
image: <numpy.ndarray shape=(2, 1, 32, 32) hash=40402718698766336 device=cpu>
mask: <numpy.ndarray shape=(2, 1, 32, 32) hash=0 device=cpu>
}
}
outputs:
batch_dict_no_nan = {
object_id: <numpy.ndarray shape=(2,) hash=14736351079544187634 device=cpu>
data: {
image: <numpy.ndarray shape=(2, 1, 32, 32) hash=40402718698766336 device=cpu>
mask: <numpy.ndarray shape=(2, 1, 32, 32) hash=0 device=cpu>
}
}
]
prepare_inputs: [
HyraxAutoencoder__prepare_inputs(batch_dict) -> batch_ndarray duration=0.00623 ms
inputs:
batch_dict = {
object_id: <numpy.ndarray shape=(2,) hash=14736351079544187634 device=cpu>
data: {
image: <numpy.ndarray shape=(2, 1, 32, 32) hash=40402718698766336 device=cpu>
mask: <numpy.ndarray shape=(2, 1, 32, 32) hash=0 device=cpu>
}
}
outputs:
batch_ndarray = <numpy.ndarray shape=(2, 1, 32, 32) hash=40402718698766336 device=cpu>
]
evaluation: [
HyraxAutoencoder__train_batch(batch) -> loss_dict duration=968 ms
inputs:
batch = <torch.Tensor shape=(2, 1, 32, 32) hash=40402718698766336 device=cuda:0>
outputs:
loss_dict = {
loss: <torch.Tensor shape=() hash=4644327510031466496 device=cuda:0>
}
HyraxAutoencoder__forward(batch) -> batch_results duration=504 ms
inputs:
batch = <torch.Tensor shape=(2, 1, 32, 32) hash=40402718698766336 device=cuda:0>
outputs:
batch_results = <torch.Tensor shape=(2, 64) hash=9301637489112383488 device=cuda:0>
]
}
The output lists five stages in pipeline order:
Stage |
What is captured |
|---|---|
|
Individual |
|
|
|
Individual |
|
|
|
Model’s |
|
Model functions such as |
The table above is printed as text from the TraceResult.__str__ implementation so that the notebook cell output and a plain print() call look identical.
Exploring stages#
You can access any stage using either attribute access or dictionary-style access — both are equivalent. Tab completion in a notebook environment will suggest the valid stage names.
[7]:
# Attribute-style access
collate_stage = trace_result.collate
print(collate_stage)
[
DataProvider__collate(batch_dicts) -> batch_dict duration=0.124 ms
inputs:
batch_dicts = <list len=2> [
{
data: {
image: <numpy.ndarray shape=(1, 32, 32) hash=66186351732588544 device=cpu>
}
object_id: '29'
}
{
data: {
image: <numpy.ndarray shape=(1, 32, 32) hash=28339009725923328 device=cpu>
}
object_id: '39'
}
]
outputs:
batch_dict = {
object_id: <numpy.ndarray shape=(2,) hash=14736351079544187634 device=cpu>
data: {
image: <numpy.ndarray shape=(2, 1, 32, 32) hash=40402718698766336 device=cpu>
mask: <numpy.ndarray shape=(2, 1, 32, 32) hash=0 device=cpu>
}
}
data__collate(samples) -> batch_dict duration=0.0376 ms
inputs:
samples = <list len=2> [
{
image: <numpy.ndarray shape=(1, 32, 32) hash=66186351732588544 device=cpu>
}
{
image: <numpy.ndarray shape=(1, 32, 32) hash=28339009725923328 device=cpu>
}
]
outputs:
batch_dict = {
image: <numpy.ndarray shape=(2, 1, 32, 32) hash=40402718698766336 device=cpu>
mask: <numpy.ndarray shape=(2, 1, 32, 32) hash=0 device=cpu>
}
DataProvider__handle_nans(batch_dict) -> batch_dict_no_nan duration=0.0665 ms
inputs:
batch_dict = {
object_id: <numpy.ndarray shape=(2,) hash=14736351079544187634 device=cpu>
data: {
image: <numpy.ndarray shape=(2, 1, 32, 32) hash=40402718698766336 device=cpu>
mask: <numpy.ndarray shape=(2, 1, 32, 32) hash=0 device=cpu>
}
}
outputs:
batch_dict_no_nan = {
object_id: <numpy.ndarray shape=(2,) hash=14736351079544187634 device=cpu>
data: {
image: <numpy.ndarray shape=(2, 1, 32, 32) hash=40402718698766336 device=cpu>
mask: <numpy.ndarray shape=(2, 1, 32, 32) hash=0 device=cpu>
}
}
]
[8]:
# Dictionary-style access — identical result
evaluation_stage = trace_result["evaluation"]
print(evaluation_stage)
[
HyraxAutoencoder__train_batch(batch) -> loss_dict duration=968 ms
inputs:
batch = <torch.Tensor shape=(2, 1, 32, 32) hash=40402718698766336 device=cuda:0>
outputs:
loss_dict = {
loss: <torch.Tensor shape=() hash=4644327510031466496 device=cuda:0>
}
HyraxAutoencoder__forward(batch) -> batch_results duration=504 ms
inputs:
batch = <torch.Tensor shape=(2, 1, 32, 32) hash=40402718698766336 device=cuda:0>
outputs:
batch_results = <torch.Tensor shape=(2, 64) hash=9301637489112383488 device=cuda:0>
]
Each stage is a TraceStage — a list of TraceCall records in the order they were executed. You can ask how many calls were captured:
[9]:
print(f"resolve_data calls : {len(trace_result.resolve_data)}")
print(f"collate calls : {len(trace_result.collate)}")
print(f"evaluation calls : {len(trace_result.evaluation)}")
resolve_data calls : 2
collate calls : 3
evaluation calls : 2
Exploring individual function calls#
Within a stage you can index calls by number (call order) or by function name. A TraceCall captures the function’s argument values and return value along with timing information.
[10]:
# Get the first call in the collate stage
first_collate_call = trace_result.collate[0]
print(first_collate_call)
DataProvider__collate(batch_dicts) -> batch_dict duration=0.124 ms
inputs:
batch_dicts = <list len=2> [
{
data: {
image: <numpy.ndarray shape=(1, 32, 32) hash=66186351732588544 device=cpu>
}
object_id: '29'
}
{
data: {
image: <numpy.ndarray shape=(1, 32, 32) hash=28339009725923328 device=cpu>
}
object_id: '39'
}
]
outputs:
batch_dict = {
object_id: <numpy.ndarray shape=(2,) hash=14736351079544187634 device=cpu>
data: {
image: <numpy.ndarray shape=(2, 1, 32, 32) hash=40402718698766336 device=cpu>
mask: <numpy.ndarray shape=(2, 1, 32, 32) hash=0 device=cpu>
}
}
[11]:
# Access the captured batch tensor directly — attribute and dict access both work
batch_dict = first_collate_call.batch_dict
print(type(batch_dict))
print(batch_dict)
<class 'dict'>
{'object_id': array(['29', '39'], dtype='<U2'), 'data': {'image': array([[[[0.06914818, 0.8157636 , 0.53917885, ..., 0.2052598 ,
0.46450907, 0.45361418],
[0.74558926, 0.5468134 , 0.79951394, ..., 0.95937926,
0.70615053, 0.7608388 ],
[0.2905739 , 0.6409456 , 0.05980945, ..., 0.2706107 ,
0.9976031 , 0.27283597],
...,
[0.60771537, 0.2282967 , 0.71244305, ..., 0.86368203,
0.64675105, 0.5531986 ],
[0.28151482, 0.34403712, 0.14771765, ..., 0.86073333,
0.01600075, 0.48290783],
[0.5655112 , 0.83720154, 0.22446716, ..., 0.40169877,
0.58955705, 0.7224778 ]]],
[[[0.04572064, 0.11397058, 0.6287335 , ..., 0.39776152,
0.6715333 , 0.3350467 ],
[0.843295 , 0.63739973, 0.28877336, ..., 0.32796383,
0.10756254, 0.92088115],
[0.203305 , 0.26796687, 0.41292202, ..., 0.66842186,
0.7503794 , 0.37287545],
...,
[0.15087664, 0.18678957, 0.6582779 , ..., 0.06698829,
0.03772938, 0.34959787],
[0.606247 , 0.9854941 , 0.45464593, ..., 0.78109044,
0.77110684, 0.2565483 ],
[0.16655219, 0.03192943, 0.08937609, ..., 0.2100066 ,
0.6990324 , 0.34367102]]]], shape=(2, 1, 32, 32), dtype=float32), 'mask': array([[[[False, False, False, ..., False, False, False],
[False, False, False, ..., False, False, False],
[False, False, False, ..., False, False, False],
...,
[False, False, False, ..., False, False, False],
[False, False, False, ..., False, False, False],
[False, False, False, ..., False, False, False]]],
[[[False, False, False, ..., False, False, False],
[False, False, False, ..., False, False, False],
[False, False, False, ..., False, False, False],
...,
[False, False, False, ..., False, False, False],
[False, False, False, ..., False, False, False],
[False, False, False, ..., False, False, False]]]],
shape=(2, 1, 32, 32))}}
[12]:
# Get a list of all calls to a particular function by providing the display name
# (visible from the print output above)
all_resolve_calls = trace_result.resolve_data["DataProvider__resolve_data"]
print(f"Number of resolve_data calls: {len(all_resolve_calls)}")
print(all_resolve_calls[0])
Number of resolve_data calls: 2
DataProvider__resolve_data(index) -> data_dict duration=0.034 ms
inputs:
index = 6
outputs:
data_dict = {
data: {
image: <numpy.ndarray shape=(1, 32, 32) hash=66186351732588544 device=cpu>
}
object_id: '29'
}
Tracing other verbs#
The trace=N keyword works the same way for infer and test:
[13]:
infer_trace = h.infer(trace=2)
print(infer_trace)
[2026-10-28 12:25:05,049 hyrax.trace:WARNING] Starting Trace
[2026-10-28 12:25:05,050 hyrax.trace:WARNING] Trace mode enabled, will only run a single batch of length 2
[2026-10-28 12:25:05,067 hyrax.models.model_registry:INFO] Setting model's self.optimizer from config: torch.optim.SGD with arguments: {'lr': 0.01, 'momentum': 0.9}.
[2026-10-28 12:25:05,068 hyrax.models.model_registry:INFO] Setting model's self.criterion from config: torch.nn.CrossEntropyLoss with default arguments.
[2026-10-28 12:25:05,068 hyrax.models.model_registry:INFO] Setting model's self.scheduler from config: torch.optim.lr_scheduler.ExponentialLR
with arguments: {'gamma': 1}.
[2026-10-28 12:25:05,068 hyrax.verbs.infer:INFO] Inference model: HyraxAutoencoder
[2026-10-28 12:25:05,069 hyrax.verbs.infer:INFO] Inference dataset(s):
{'infer': Name: data (primary dataset)
Dataset class: HyraxRandomDataset
Data location: /home/drew/code/hyrax/docs/pre_executed/trace_data
Selected items: 50
Primary ID field: object_id
Requested fields: image
}
2026-10-28 12:25:05,070 ignite.distributed.auto.auto_dataloader INFO: Use data loader kwargs for dataset '<torch.utils.data.da':
{'sampler': None, 'batch_size': 2, 'collate_fn': <bound method CollationMixin.collate of Name: data (primary dataset)
Dataset class: HyraxRandomDataset
Data location: /home/drew/code/hyrax/docs/pre_executed/trace_data
Selected items: 50
Primary ID field: object_id
Requested fields: image
>, 'pin_memory': True}
[2026-10-28 12:25:05,080 hyrax.models.model_utils:INFO] Updated config['infer']['model_weights_file'] to: /home/drew/code/hyrax/docs/pre_executed/results/20261028-122502-train-tWy2/example_model.pth
[2026-10-28 12:25:05,081 hyrax.verbs.infer:INFO] Saving inference results at: /home/drew/code/hyrax/docs/pre_executed/results/20261028-122505-infer-umzD
[2026-10-28 12:25:05,103 hyrax.pytorch_ignite:INFO] Total evaluation time: 0.02[s]
[2026-10-28 12:25:05,104 hyrax.datasets.result_dataset:INFO] Optimizing Lance table after 1 batches
[2026-10-28 12:25:05,112 hyrax.datasets.result_dataset:INFO] Lance table optimization complete
[2026-10-28 12:25:05,117 hyrax.verbs.infer:INFO] Inference Complete.
Trace Stages {
dataset_getter: [
data__get_image(index) -> image duration=0.00595 ms
inputs:
index = 0
outputs:
image = <numpy.ndarray shape=(1, 32, 32) hash=26660822658842624 device=cpu>
data__get_object_id(index) -> object_id duration=0.00508 ms
inputs:
index = 0
outputs:
object_id = '23'
data__get_image(index) -> image duration=0.000414 ms
inputs:
index = 1
outputs:
image = <numpy.ndarray shape=(1, 32, 32) hash=7316936887107584 device=cpu>
data__get_object_id(index) -> object_id duration=0.000404 ms
inputs:
index = 1
outputs:
object_id = '24'
]
resolve_data: [
DataProvider__resolve_data(index) -> data_dict duration=0.0248 ms
inputs:
index = 0
outputs:
data_dict = {
data: {
image: <numpy.ndarray shape=(1, 32, 32) hash=26660822658842624 device=cpu>
}
object_id: '23'
}
DataProvider__resolve_data(index) -> data_dict duration=0.00515 ms
inputs:
index = 1
outputs:
data_dict = {
data: {
image: <numpy.ndarray shape=(1, 32, 32) hash=7316936887107584 device=cpu>
}
object_id: '24'
}
]
field_level_collation: []
collate: [
DataProvider__collate(batch_dicts) -> batch_dict duration=0.11 ms
inputs:
batch_dicts = <list len=2> [
{
data: {
image: <numpy.ndarray shape=(1, 32, 32) hash=26660822658842624 device=cpu>
}
object_id: '23'
}
{
data: {
image: <numpy.ndarray shape=(1, 32, 32) hash=7316936887107584 device=cpu>
}
object_id: '24'
}
]
outputs:
batch_dict = {
object_id: <numpy.ndarray shape=(2,) hash=9290790550866539747 device=cpu>
data: {
image: <numpy.ndarray shape=(2, 1, 32, 32) hash=20065364041793536 device=cpu>
mask: <numpy.ndarray shape=(2, 1, 32, 32) hash=0 device=cpu>
}
}
data__collate(samples) -> batch_dict duration=0.0362 ms
inputs:
samples = <list len=2> [
{
image: <numpy.ndarray shape=(1, 32, 32) hash=26660822658842624 device=cpu>
}
{
image: <numpy.ndarray shape=(1, 32, 32) hash=7316936887107584 device=cpu>
}
]
outputs:
batch_dict = {
image: <numpy.ndarray shape=(2, 1, 32, 32) hash=20065364041793536 device=cpu>
mask: <numpy.ndarray shape=(2, 1, 32, 32) hash=0 device=cpu>
}
DataProvider__handle_nans(batch_dict) -> batch_dict_no_nan duration=0.0523 ms
inputs:
batch_dict = {
object_id: <numpy.ndarray shape=(2,) hash=9290790550866539747 device=cpu>
data: {
image: <numpy.ndarray shape=(2, 1, 32, 32) hash=20065364041793536 device=cpu>
mask: <numpy.ndarray shape=(2, 1, 32, 32) hash=0 device=cpu>
}
}
outputs:
batch_dict_no_nan = {
object_id: <numpy.ndarray shape=(2,) hash=9290790550866539747 device=cpu>
data: {
image: <numpy.ndarray shape=(2, 1, 32, 32) hash=20065364041793536 device=cpu>
mask: <numpy.ndarray shape=(2, 1, 32, 32) hash=0 device=cpu>
}
}
]
prepare_inputs: [
HyraxAutoencoder_inst_prepare_inputs(batch_dict) -> batch_ndarray duration=0.00256 ms
inputs:
batch_dict = {
object_id: <numpy.ndarray shape=(2,) hash=9290790550866539747 device=cpu>
data: {
image: <numpy.ndarray shape=(2, 1, 32, 32) hash=20065364041793536 device=cpu>
mask: <numpy.ndarray shape=(2, 1, 32, 32) hash=0 device=cpu>
}
}
outputs:
batch_ndarray = <numpy.ndarray shape=(2, 1, 32, 32) hash=20065364041793536 device=cpu>
]
evaluation: [
HyraxAutoencoder__infer_batch(batch) -> batch_results duration=0.525 ms
inputs:
batch = <torch.Tensor shape=(2, 1, 32, 32) hash=20065364041793536 device=cuda:0>
outputs:
batch_results = <torch.Tensor shape=(2, 64) hash=57208774867288064 device=cuda:0>
HyraxAutoencoder__forward(batch) -> batch_results duration=0.519 ms
inputs:
batch = <torch.Tensor shape=(2, 1, 32, 32) hash=20065364041793536 device=cuda:0>
outputs:
batch_results = <torch.Tensor shape=(2, 64) hash=57208774867288064 device=cuda:0>
]
}
Instrumenting custom models and datasets#
If you have written a custom model or dataset class you can opt specific methods into trace capture using the @trace_model_func and @trace_dataset_func decorators.
from hyrax.trace import trace_model_func, trace_dataset_func
class MyModel(nn.Module):
@trace_model_func
def my_custom_forward(self, batch):
# This call will appear in the 'evaluation' stage of the TraceResult
...
class MyDataset(HyraxDataset):
@trace_dataset_func
def get_image(self, index):
# This call will appear in the 'dataset_getter' stage of the TraceResult
...
The decorators add a small overhead to every call, so they are intended for use during development and debugging rather than in production. Remove them (or use them selectively) once you are satisfied with how your data is flowing.