ONNX export update

This commit is contained in:
Glenn Jocher 2020-04-15 12:12:59 -07:00
parent 628028c617
commit b8c3644a18
1 changed files with 2 additions and 2 deletions

View File

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