F.pad 的使用参考 Pytorch之torch.nn.functional.pad函数详解 。
F.pad 可用在通道扩充上,在动态batch输入时,导出的onnx非常简洁。
# 简化的示例代码 import torch import torch.nn.functional as F x = torch.randn(1,3,5,4) added_chan_num = 2 use_pad = True if (not use_pad): z = torch.zeros(1,added_chan_num,5,4) x_a = torch.cat((x,z),dim=1) else: x_a = F.pad(x, (0,0,0,0,0,added_chan_num), 'constant', 0) print(x_a)
onnx效果:
欢迎分享,转载请注明来源:内存溢出
评论列表(0条)