Masaki Yano
commited on
Commit
·
71c1fa1
1
Parent(s):
583f89f
chore(tools): add decode_in_inference option to export_torchscript (#1442)
Browse files
tools/export_torchscript.py
CHANGED
@@ -27,6 +27,11 @@ def make_parser():
|
|
27 |
parser.add_argument("-expn", "--experiment-name", type=str, default=None)
|
28 |
parser.add_argument("-n", "--name", type=str, default=None, help="model name")
|
29 |
parser.add_argument("-c", "--ckpt", default=None, type=str, help="ckpt path")
|
|
|
|
|
|
|
|
|
|
|
30 |
parser.add_argument(
|
31 |
"opts",
|
32 |
help="Modify config options using the command-line",
|
@@ -61,7 +66,7 @@ def main():
|
|
61 |
if "model" in ckpt:
|
62 |
ckpt = ckpt["model"]
|
63 |
model.load_state_dict(ckpt)
|
64 |
-
model.head.decode_in_inference =
|
65 |
|
66 |
logger.info("loading checkpoint done.")
|
67 |
dummy_input = torch.randn(args.batch_size, 3, exp.test_size[0], exp.test_size[1])
|
|
|
27 |
parser.add_argument("-expn", "--experiment-name", type=str, default=None)
|
28 |
parser.add_argument("-n", "--name", type=str, default=None, help="model name")
|
29 |
parser.add_argument("-c", "--ckpt", default=None, type=str, help="ckpt path")
|
30 |
+
parser.add_argument(
|
31 |
+
"--decode_in_inference",
|
32 |
+
action="store_true",
|
33 |
+
help="decode in inference or not"
|
34 |
+
)
|
35 |
parser.add_argument(
|
36 |
"opts",
|
37 |
help="Modify config options using the command-line",
|
|
|
66 |
if "model" in ckpt:
|
67 |
ckpt = ckpt["model"]
|
68 |
model.load_state_dict(ckpt)
|
69 |
+
model.head.decode_in_inference = args.decode_in_inference
|
70 |
|
71 |
logger.info("loading checkpoint done.")
|
72 |
dummy_input = torch.randn(args.batch_size, 3, exp.test_size[0], exp.test_size[1])
|