pytorch Flatten展平

经过测试,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)

版权声明:本文为jacke121原创文章,遵循CC 4.0 BY-SA版权协议,转载请附上原文出处链接和本声明。