def setUp(self):
TikTorch.read_config = lambda self: self
tiktorch = TikTorch(build_directory=".")
if self.MODEL == "HEAVY":
tiktorch._model = nn.Sequential(
nn.Conv2d(1, 512, 9, padding=4),
nn.ELU(),
nn.Conv2d(512, 512, 9, padding=4),
nn.ELU(),
nn.Conv2d(512, 512, 9, padding=4),
nn.ELU(),
nn.Conv2d(512, 512, 9, padding=4),
nn.ELU(),
nn.Conv2d(512, 1, 9, padding=4),
)
elif self.MODEL == "LITE":
tiktorch._model = nn.Conv2d(1, 1, 1)
tiktorch._config = {"input_shape": [1, 512, 512], "dynamic_input_shape": "(32 * (nH + 1), 32 * (nW + 1))"}
tiktorch._set_handler(tiktorch._model)
self.tiktorch = tiktorch
def setUp(self):
TikTorch.read_config = lambda self: self
tiktorch = TikTorch(build_directory='.')
if self.MODEL == 'HEAVY':
tiktorch._model = nn.Sequential(
nn.Conv2d(1, 512, 9, padding=4),
nn.ELU(),
nn.Conv2d(512, 512, 9, padding=4),
nn.ELU(),
nn.Conv2d(512, 512, 9, padding=4),
nn.ELU(),
nn.Conv2d(512, 512, 9, padding=4),
nn.ELU(),
nn.Conv2d(512, 1, 9, padding=4),
)
elif self.MODEL == 'LITE':
tiktorch._model = nn.Conv2d(1, 1, 1)
tiktorch._config = {
'input_shape': [1, 512, 512],
'dynamic_input_shape': '(32 * (nH + 1), 32 * (nW + 1))'
}
tiktorch._set_handler(tiktorch._model)
self.tiktorch = tiktorch