通過測試,cpu模式中,用Module的方式比直接在推理中展平平均快1~5ms python
from torch import nn import torch import math class Flatten(Module): def forward(self, input): return input.view(input.size(0), -1)
直接展平: 測試
out = self.avgpool(out) # out = self.conv_6_flatten(out) out = out.view(out.size(0), -1)