From 84b4b6229f871ad02efaf8b6afd32c804ad13829 Mon Sep 17 00:00:00 2001 From: Yonghye Kwon Date: Tue, 11 May 2021 11:10:01 +0900 Subject: [PATCH 1/2] cleanup for flatten layer use official flatten function --- hardnet.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/hardnet.py b/hardnet.py index cdb7a13..ae29bde 100644 --- a/hardnet.py +++ b/hardnet.py @@ -7,9 +7,7 @@ class Flatten(nn.Module): def __init__(self): super().__init__() def forward(self, x): - return x.view(x.data.size(0),-1) - - + return x.flatten(start_dim=1) class CombConvLayer(nn.Sequential): def __init__(self, in_channels, out_channels, kernel=1, stride=1, dropout=0.1, bias=False): From 2ea54842906fa0d41fd3f47b762d40dba30a348e Mon Sep 17 00:00:00 2001 From: Yonghye Kwon Date: Thu, 13 May 2021 16:56:29 +0900 Subject: [PATCH 2/2] Update hardnet.py --- hardnet.py | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/hardnet.py b/hardnet.py index ae29bde..9f93b5d 100644 --- a/hardnet.py +++ b/hardnet.py @@ -3,12 +3,6 @@ import torch.nn as nn import torch.nn.functional as F -class Flatten(nn.Module): - def __init__(self): - super().__init__() - def forward(self, x): - return x.flatten(start_dim=1) - class CombConvLayer(nn.Sequential): def __init__(self, in_channels, out_channels, kernel=1, stride=1, dropout=0.1, bias=False): super().__init__() @@ -194,7 +188,7 @@ def __init__(self, depth_wise=False, arch=85, pretrained=True, weight_path=''): self.base.append ( nn.Sequential( nn.AdaptiveAvgPool2d((1,1)), - Flatten(), + nn.Flatten(), nn.Dropout(drop_rate), nn.Linear(ch, 1000) ))