经过测试,cpu模式中,用Module的方式比直接在推理中展平平均快1~5ms
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)
展平为向量
label = label.view(-1)