Update eval_onnx.py
Browse files- eval_onnx.py +3 -5
eval_onnx.py
CHANGED
@@ -21,8 +21,6 @@ parser.add_argument(
|
|
21 |
"--data_dir",
|
22 |
default="/workspace/dataset/imagenet",
|
23 |
help="Directory of dataset")
|
24 |
-
parser.add_argument(
|
25 |
-
"--batch_size", default=1, type=int, help="Evaluation batch size")
|
26 |
parser.add_argument(
|
27 |
"--ipu",
|
28 |
action="store_true",
|
@@ -87,7 +85,7 @@ def accuracy(output: torch.Tensor,
|
|
87 |
return res
|
88 |
|
89 |
def prepare_data_loader(data_dir: str,
|
90 |
-
batch_size: int =
|
91 |
workers: int = 8) -> torch.utils.data.DataLoader:
|
92 |
"""Returns a validation data loader of ImageNet by given `data_dir`.
|
93 |
Args:
|
@@ -132,7 +130,7 @@ def val_imagenet():
|
|
132 |
ort_session = onnxruntime.InferenceSession(
|
133 |
args.onnx_model, providers=providers, provider_options=provider_options)
|
134 |
|
135 |
-
val_loader = prepare_data_loader(args.data_dir
|
136 |
|
137 |
top1 = AverageMeter('Acc@1', ':6.2f')
|
138 |
top5 = AverageMeter('Acc@5', ':6.2f')
|
@@ -158,4 +156,4 @@ def val_imagenet():
|
|
158 |
return top1.avg, top5.avg
|
159 |
|
160 |
if __name__ == '__main__':
|
161 |
-
val_imagenet()
|
|
|
21 |
"--data_dir",
|
22 |
default="/workspace/dataset/imagenet",
|
23 |
help="Directory of dataset")
|
|
|
|
|
24 |
parser.add_argument(
|
25 |
"--ipu",
|
26 |
action="store_true",
|
|
|
85 |
return res
|
86 |
|
87 |
def prepare_data_loader(data_dir: str,
|
88 |
+
batch_size: int = 1,
|
89 |
workers: int = 8) -> torch.utils.data.DataLoader:
|
90 |
"""Returns a validation data loader of ImageNet by given `data_dir`.
|
91 |
Args:
|
|
|
130 |
ort_session = onnxruntime.InferenceSession(
|
131 |
args.onnx_model, providers=providers, provider_options=provider_options)
|
132 |
|
133 |
+
val_loader = prepare_data_loader(args.data_dir)
|
134 |
|
135 |
top1 = AverageMeter('Acc@1', ':6.2f')
|
136 |
top5 = AverageMeter('Acc@5', ':6.2f')
|
|
|
156 |
return top1.avg, top5.avg
|
157 |
|
158 |
if __name__ == '__main__':
|
159 |
+
val_imagenet()
|