Browse Source

修改yolox白盒水印处理方式

liyan 10 months ago
parent
commit
07e23e562e
1 changed files with 4 additions and 2 deletions
  1. 4 2
      watermark_generate/deals/yolox_pytorch_white_embed.py

+ 4 - 2
watermark_generate/deals/yolox_pytorch_white_embed.py

@@ -40,6 +40,7 @@ import torch.nn.functional as F
 """
 import torch.nn.functional as F
 import os
+import numpy as np
 """
     # 文件替换
     modify_file.replace_block_in_file(project_file, old_source_block, new_source_block)
@@ -180,9 +181,10 @@ class ModelEncoder:
         return [int(x) for x in binary_representation]
 
     def save_tensor(self, tensor, save_path):
-        assert save_path.endswith('.pt') or save_path.endswith('.pth'), "权重保存文件必须以.pt或.pth结尾"
         os.makedirs(os.path.dirname(save_path), exist_ok=True)
-        torch.save(tensor, save_path)
+        tensor = tensor.cpu()
+        numpy_array = tensor.numpy()
+        np.save(save_path, numpy_array)
 
     def flatten_parameters(self, weights):
         return torch.cat([torch.mean(x, dim=3).reshape(-1)