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])
|