解決pytorch resnet50模型導出onnx時動態batch_size難題
本文介紹如何將基于ResNet50的PyTorch模型導出為ONNX格式，重點解決動態batch_size導致的導出問題。原始代碼中，imageretrievalnet類和gem類存在一些與ONNX導出不兼容的因素，主要包括gem類中可學習參數self.p以及imageretrievalnet類中未使用的self.lwhiten屬性。這些動態元素阻礙了ONNX的shape推斷，導致導出失敗。
為了解決這個問題，我們需要修改這兩個類以適應ONNX導出流程。具體修改如下：
首先，修改gem類，將self.p參數改為直接賦值的常量，不再作為可學習參數：
class gem(nn.Module): def __init__(self, p=3, eps=1e-6): super(gem, self).__init__() self.p = p # 直接賦值常量值 self.eps = eps def forward(self, x: torch.Tensor) -> torch.Tensor: return gem_op(x, p=self.p, eps=self.eps) # 使用自定義的gem_op函數，避免直接使用類名調用
然后，簡化imageretrievalnet類，去除未使用的self.lwhiten屬性：
class imageretrievalnet(nn.Module): def __init__(self, dim: int = 512): super(imageretrievalnet, self).__init__() resnet50_model = models.resnet50() features = list(resnet50_model.children())[:-2] self.features = nn.Sequential(*features) self.pool = gem() self.whiten = nn.Linear(2048, dim, bias=True) # 使用nn.Linear self.norm = l2n() def forward(self, x: torch.Tensor): o: torch.Tensor = self.features(x) pooled_t = self.pool(o) normed_t: torch.Tensor = self.norm(pooled_t) o: torch.Tensor = normed_t.squeeze(-1).squeeze(-1) if self.whiten is not None: whitened_t = self.whiten(o) normed_t: torch.Tensor = self.norm(whitened_t) o = normed_t return o.permute(1, 0)
通過以上修改，消除了動態參數帶來的不確定性，使ONNX導出能夠順利進行。 使用修改后的imageretrievalnet類，并利用torch.onnx.export函數，指定dynamic_axes參數處理動態batch_size，即可成功導出ONNX模型：
model = imageretrievalnet() batch_size = 4 input_shape = (batch_size, 3, 224, 224) input_data = torch.randn(input_shape) torch.onnx.export( model, input_data, "resnet50.onnx", input_names=["input"], output_names=["output"], opset_version=12, dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}} )
記住根據實際情況調整opset_version參數。 通過這些修改，即可成功導出支持動態batch_size的ResNet50 ONNX模型。 請注意，代碼中添加了gem_op函數的假設，該函數應該實現gem類的功能，以避免在ONNX導出過程中直接使用類名調用。
? 版權聲明
文章版權歸作者所有，未經允許請勿轉載。
THE END
