ONNX export update
This commit is contained in:
parent
628028c617
commit
b8c3644a18
|
@ -227,7 +227,7 @@ class Darknet(nn.Module):
|
|||
# Darknet Header https://github.com/AlexeyAB/darknet/issues/2914#issuecomment-496675346
|
||||
self.version = np.array([0, 2, 5], dtype=np.int32) # (int32) version info: major, minor, revision
|
||||
self.seen = np.array([0], dtype=np.int64) # (int64) number of images seen during training
|
||||
self.info(verbose) # print model description
|
||||
self.info(verbose) if not ONNX_EXPORT else None # print model description
|
||||
|
||||
def forward(self, x, augment=False, verbose=False):
|
||||
|
||||
|
@ -324,7 +324,7 @@ class Darknet(nn.Module):
|
|||
break
|
||||
fused_list.append(a)
|
||||
self.module_list = fused_list
|
||||
self.info() # yolov3-spp reduced from 225 to 152 layers
|
||||
self.info() if not ONNX_EXPORT else None # yolov3-spp reduced from 225 to 152 layers
|
||||
|
||||
def info(self, verbose=False):
|
||||
torch_utils.model_info(self, verbose)
|
||||
|
|
Loading…
Reference in New Issue