提交 d6817995 authored 作者: Glenn Jocher's avatar Glenn Jocher

Update export.py

上级 6f3db5e6
...@@ -9,7 +9,7 @@ import argparse ...@@ -9,7 +9,7 @@ import argparse
import torch import torch
import torch.nn as nn import torch.nn as nn
from models.common import Conv import models
from models.experimental import attempt_load from models.experimental import attempt_load
from utils.activations import Hardswish from utils.activations import Hardswish
from utils.general import set_logging from utils.general import set_logging
...@@ -33,10 +33,10 @@ if __name__ == '__main__': ...@@ -33,10 +33,10 @@ if __name__ == '__main__':
# Update model # Update model
for k, m in model.named_modules(): for k, m in model.named_modules():
m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatability m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatability
if isinstance(m, Conv) and isinstance(m.act, nn.Hardswish): if isinstance(m, models.common.Conv) and isinstance(m.act, nn.Hardswish):
m.act = Hardswish() # assign activation m.act = Hardswish() # assign activation
# if isinstance(m, Detect): # if isinstance(m, models.yolo.Detect):
# m.forward = m.forward_export # assign forward (optional) # m.forward = m.forward_export # assign forward (optional)
model.model[-1].export = True # set Detect() layer export=True model.model[-1].export = True # set Detect() layer export=True
y = model(img) # dry run y = model(img) # dry run
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论