output_featuremaps_only -> return_featuremaps

pull/218/head
KaiyangZhou 2019-08-05 14:01:39 +01:00
parent 0224a5d433
commit a20f8a66bd
2 changed files with 6 additions and 6 deletions

View File

@ -324,11 +324,11 @@ class Engine(object):
# forward to get convolutional feature maps
try:
outputs = self.model(imgs, output_featuremaps_only=True)
outputs = self.model(imgs, return_featuremaps=True)
except TypeError:
raise TypeError('forward() got unexpected keyword argument "output_featuremaps_only". ' \
'Please add output_featuremaps_only as an input argument to forward(). When ' \
'output_featuremaps_only=True, return feature maps only.')
raise TypeError('forward() got unexpected keyword argument "return_featuremaps". ' \
'Please add return_featuremaps as an input argument to forward(). When ' \
'return_featuremaps=True, return feature maps only.')
if outputs.dim() != 4:
raise ValueError('The model output is supposed to have ' \

View File

@ -292,9 +292,9 @@ class OSNet(nn.Module):
x = self.conv5(x)
return x
def forward(self, x, output_featuremaps_only=False):
def forward(self, x, return_featuremaps=False):
x = self.featuremaps(x)
if output_featuremaps_only:
if return_featuremaps:
return x
v = self.global_avgpool(x)
v = v.view(v.size(0), -1)