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
|
# 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.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.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):
|
def forward(self, x, augment=False, verbose=False):
|
||||||
|
|
||||||
|
@ -324,7 +324,7 @@ class Darknet(nn.Module):
|
||||||
break
|
break
|
||||||
fused_list.append(a)
|
fused_list.append(a)
|
||||||
self.module_list = fused_list
|
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):
|
def info(self, verbose=False):
|
||||||
torch_utils.model_info(self, verbose)
|
torch_utils.model_info(self, verbose)
|
||||||
|
|
Loading…
Reference in New Issue