소스 검색

修改白盒水印嵌入代码,调整pytorch的坐标顺序与tensorflow一致

liyan 10 달 전
부모
커밋
5a7e1f6473
1개의 변경된 파일1개의 추가작업 그리고 0개의 파일을 삭제
  1. 1 0
      watermark_generate/deals/yolox_pytorch_white_embed.py

+ 1 - 0
watermark_generate/deals/yolox_pytorch_white_embed.py

@@ -187,6 +187,7 @@ class ModelEncoder:
         np.save(save_path, numpy_array)
 
     def flatten_parameters(self, weights):
+        weights = [weight.permute(2, 3, 1, 0) for weight in weights]
         return torch.cat([torch.mean(x, dim=3).reshape(-1)
                           for x in weights])