def __init__(self, training_name, data_path, training_results_path):
super().__init__(training_name, data_path, training_results_path)
# Config of the data
self.data_dataset = FashionMNISTDataset
# Config of the model
self.model_model = lambda config: Module.create_from_file(
"deeptech/examples/mnist_model.json",
"MNISTModel",
num_classes=10,
logits=True)
# Config for training
self.training_loss = SparseCrossEntropyLossFromLogits
self.training_optimizer = smart_optimizer(SGD)
self.training_trainer = SupervisedTrainer
self.training_epochs = 10
self.training_batch_size = 32
def setUp(self) -> None:
self.module = Module.create_from_file("tests/json_nets/vgg16_bn.jsonc", "VGG16_bn", logits=True)
self.input_data = torch.from_numpy(np.zeros((1, 128, 128, 3), dtype=np.float32))
self.result = self.module(self.input_data)