add model eval and spec

examples/mask_rcnn.py

@@ -183,7 +183,7 @@ class Masker(object):
 
 masker = Masker(threshold=0.5, padding=1)
 
-def compute_prediction(original_image, model_type='tiny'):
+def compute_prediction(original_image, model):
   # apply pre-processing to image
   image = transforms(original_image).numpy()
   image = Tensor(image, requires_grad=False)
@@ -200,11 +200,19 @@ def compute_prediction(original_image, model_type='tiny'):
     masks = prediction.get_field("mask")
     # always single image is passed at a time
     masks = masker([masks], [prediction])[0]
-    if model_type != 'tiny':
-      masks = torch.tensor(masks.numpy())
     prediction.add_field("mask", masks)
   return prediction
 
+def compute_prediction_batched(batch, model):
+  # apply pre-processing to image
+  imgs = []
+  for img in batch:
+    imgs.append(transforms(img).numpy())
+  image = [Tensor(image, requires_grad=False) for image in imgs]
+  predictions = model(image)
+  del image
+  return predictions
+
 palette = torch.tensor([2 ** 25 - 1, 2 ** 15 - 1, 2 ** 21 - 1])
 
 def findContours(*args, **kwargs):
@@ -342,7 +350,7 @@ if __name__ == '__main__':
   model_tiny = MaskRCNN(resnet)
   model_tiny.load_from_pretrained()
   img = Image.open(args.image)
-  result = compute_prediction(img)
+  result = compute_prediction(img, model_tiny)
   top_result_tiny = select_top_predictions(result, confidence_threshold=args.threshold)
   bbox_image = overlay_boxes(img, top_result_tiny)
   mask_image = overlay_mask(bbox_image, top_result_tiny)

examples/mlperf/model_eval.py

@@ -184,12 +184,49 @@ def eval_bert():
 
     st = time.perf_counter()
 
+def eval_mrcnn():
+  from tqdm import tqdm
+  from models.mask_rcnn import MaskRCNN
+  from models.resnet import ResNet
+  from datasets.coco import BASEDIR, images, convert_prediction_to_coco_bbox, convert_prediction_to_coco_mask, accumulate_predictions_for_coco, evaluate_predictions_on_coco, iterate
+  from examples.mask_rcnn import compute_prediction_batched, Image
+  mdl = MaskRCNN(ResNet(50, num_classes=None, stride_in_1x1=True))
+  mdl.load_from_pretrained()
+
+  bbox_output = '/tmp/results_bbox.json'
+  mask_output = '/tmp/results_mask.json'
+
+  accumulate_predictions_for_coco([], bbox_output, rm=True)
+  accumulate_predictions_for_coco([], mask_output, rm=True)
+
+  #TODO: bs > 1 not as accurate 
+  bs = 1 
+
+  for batch in tqdm(iterate(images, bs=bs), total=len(images)//bs):
+    batch_imgs = []
+    for image_row in batch:
+      image_name = image_row['file_name']
+      img = Image.open(BASEDIR/f'val2017/{image_name}')
+      batch_imgs.append(img)
+    batch_result = compute_prediction_batched(batch_imgs, mdl)
+    for image_row, result in zip(batch, batch_result):
+      image_name = image_row['file_name']
+      box_pred = convert_prediction_to_coco_bbox(image_name, result)
+      mask_pred = convert_prediction_to_coco_mask(image_name, result)
+      accumulate_predictions_for_coco(box_pred, bbox_output)
+      accumulate_predictions_for_coco(mask_pred, mask_output)
+    del batch_imgs
+    del batch_result
+
+  evaluate_predictions_on_coco(bbox_output, iou_type='bbox')
+  evaluate_predictions_on_coco(bbox_output, iou_type='segm')
+
 if __name__ == "__main__":
   # inference only
   Tensor.training = False
   Tensor.no_grad = True
 
-  models = getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert").split(",")
+  models = getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert,mrcnn").split(",")
   for m in models:
     nm = f"eval_{m}"
     if nm in globals():

examples/mlperf/model_spec.py

@@ -5,7 +5,8 @@ import numpy as np
 
 def test_model(model, *inputs):
   GlobalCounters.reset()
-  model(*inputs).numpy()
+  out = model(*inputs)
+  if isinstance(out, Tensor): out = out.numpy()
   # TODO: return event future to still get the time_sum_s without DEBUG=2
   print(f"{GlobalCounters.global_ops*1e-9:.2f} GOPS, {GlobalCounters.time_sum_s*1000:.2f} ms")
 
@@ -49,12 +50,19 @@ def spec_bert():
   tt = Tensor(np.random.randint(0, 2, (1, 384)).astype(np.float32))
   test_model(mdl, x, am, tt)
 
+def spec_mrcnn():
+  from models.mask_rcnn import MaskRCNN, ResNet
+  mdl = MaskRCNN(ResNet(50, num_classes=None, stride_in_1x1=True))
+  mdl.load_from_pretrained()
+  x = Tensor.randn(3, 224, 224)
+  test_model(mdl, [x])
+
 if __name__ == "__main__":
   # inference only for now
   Tensor.training = False
   Tensor.no_grad = True
 
-  for m in getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert").split(","):
+  for m in getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert,mrcnn").split(","):
     nm = f"spec_{m}"
     if nm in globals():
       print(f"testing {m}")