def setup_module(module):
global input_batch, vision_model, managed_model, single_output, batch_output
with open(os.path.join(DIR_PATH, "cat.jpg"), 'rb') as f:
image_bytes = f.read()
input_batch = [image_bytes]
vision_model = VisionModel(device=device)
single_output = vision_model.batch_prediction(input_batch)
batch_output = vision_model.batch_prediction(input_batch * BATCH_SIZE)
managed_model = ManagedVisionModel()
managed_model.init_model()
class ManagedVisionModel(ManagedModel):
def init_model(self):
self.model = VisionModel(device=device)
def predict(self, batch):
return self.model.batch_prediction(batch)
def setup_module(module):
global input_batch, vision_model, managed_model
device = "cpu" # in case ci environment do not have gpu
with open(os.path.join(DIR_PATH, "cat.jpg"), 'rb') as f:
image_bytes = f.read()
input_batch = [image_bytes]
vision_model = VisionModel(device=device)
class ManagedVisionModel(ManagedModel):
def init_model(self):
self.model = VisionModel(device=device)
def predict(self, batch):
return self.model.batch_prediction(batch)
managed_model = ManagedVisionModel()
managed_model.init_model()
def init_model(self):
self.model = VisionModel(device=device)