瀏覽代碼

Consistent saved_model output format (#7032)

modifyDataloader
Mrinal Jain GitHub 2 年之前
父節點
當前提交
4effd064b1
沒有發現已知的金鑰在資料庫的簽署中 GPG Key ID: 4AEE18F83AFDEB23
共有 2 個文件被更改,包括 2 次插入2 次删除
  1. +1
    -1
      export.py
  2. +1
    -1
      models/common.py

+ 1
- 1
export.py 查看文件

@@ -275,7 +275,7 @@ def export_saved_model(model, im, file, dynamic,
m = m.get_concrete_function(spec)
frozen_func = convert_variables_to_constants_v2(m)
tfm = tf.Module()
tfm.__call__ = tf.function(lambda x: frozen_func(x), [spec])
tfm.__call__ = tf.function(lambda x: frozen_func(x)[0], [spec])
tfm.__call__(im)
tf.saved_model.save(
tfm,

+ 1
- 1
models/common.py 查看文件

@@ -441,7 +441,7 @@ class DetectMultiBackend(nn.Module):
else: # TensorFlow (SavedModel, GraphDef, Lite, Edge TPU)
im = im.permute(0, 2, 3, 1).cpu().numpy() # torch BCHW to numpy BHWC shape(1,320,192,3)
if self.saved_model: # SavedModel
y = (self.model(im, training=False) if self.keras else self.model(im)[0]).numpy()
y = (self.model(im, training=False) if self.keras else self.model(im)).numpy()
elif self.pb: # GraphDef
y = self.frozen_func(x=self.tf.constant(im)).numpy()
else: # Lite or Edge TPU

Loading…
取消
儲存