def _build_loaders(self):
train_path = self.get("data/paths/train", ensure_exists=True)
validate_path = self.get("data/paths/validate", ensure_exists=True)
self.train_loader = get_dataloader(path=train_path,
**self.get("data/loader_kwargs",
ensure_exists=True))
self.validate_loader = get_dataloader(path=validate_path,
**self.get("data/loader_kwargs",
ensure_exists=True))
def test_loader(self):
from loader import get_dataloader
path = self.DATASET_PATH
batch_size = 5
dataloader = get_dataloader(
batch_size=batch_size, shuffle=False, num_workers=0, path=path
)
batch = next(iter(dataloader))
self.assertEqual(len(batch), self.NUM_KEYS_IN_BATCH)
# Testing that all the keys in the batch have the batch_size
keys_in_batch = list(batch.keys())
for key in keys_in_batch:
self.assertEqual(len(batch[key]), batch_size)
dataloader = get_dataloader(
batch_size=batch_size, shuffle=False, num_workers=0, path=[path, path]
)
batch = next(iter(dataloader))
import onnx
from onnx_tf.backend import prepare
from models import ContactTracingTransformer
from loader import get_dataloader
import torch
path = "output.pkl"
dataloader = get_dataloader(batch_size=1,
shuffle=False,
num_workers=0,
path=path)
batch = next(iter(dataloader))
#Load ONNX model
onnx_model = onnx.load("model_onnx_10.onnx")
tf_model = prepare(onnx_model)
#Inputs to the model
print('inputs:', tf_model.inputs)
# Output nodes from the model
print('outputs:', tf_model.outputs)
# All nodes in the model
# print('tensor_dict:')
# print(tf_model.tensor_dict)
output = tf_model.run(batch)
print(output)
tf_model.export_graph('tf_graph2.pb')
#Sanity check with the PyTorch Model
ctt = ContactTracingTransformer(pool_latent_entities=False,
use_logit_sink=False)
ctt.load_state_dict(torch.load('model.pth'))