diff --git a/examples/mask_rcnn.py b/examples/mask_rcnn.py index 604aa0ad92..3b1b37adf7 100644 --- a/examples/mask_rcnn.py +++ b/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) diff --git a/examples/mlperf/model_eval.py b/examples/mlperf/model_eval.py index 0c6b925447..cc5d610dfb 100644 --- a/examples/mlperf/model_eval.py +++ b/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(): diff --git a/examples/mlperf/model_spec.py b/examples/mlperf/model_spec.py index 69ff7caec1..89ac387286 100644 --- a/examples/mlperf/model_spec.py +++ b/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}")