Update export.py
parent
6f3db5e662
commit
d68179953d
|
@ -9,7 +9,7 @@ import argparse
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from models.common import Conv
|
||||
import models
|
||||
from models.experimental import attempt_load
|
||||
from utils.activations import Hardswish
|
||||
from utils.general import set_logging
|
||||
|
@ -33,10 +33,10 @@ if __name__ == '__main__':
|
|||
# Update model
|
||||
for k, m in model.named_modules():
|
||||
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
|
||||
# if isinstance(m, Detect):
|
||||
# m.forward = m.forward_export # assign forward (optional)
|
||||
# if isinstance(m, models.yolo.Detect):
|
||||
# m.forward = m.forward_export # assign forward (optional)
|
||||
model.model[-1].export = True # set Detect() layer export=True
|
||||
y = model(img) # dry run
|
||||
|
||||
|
|
Loading…
Reference in New Issue