def get_out_shape(layer, input_shape):
if hasattr(layer, 'out_channels'):
return get_output_shape(input_shape,
kernel_size=layer.kernel_size,
stride=layer.stride,
padding=layer.padding,
dilation=layer.dilation)
elif hasattr(layer, 'out_features'):
return []
elif hasattr(layer, 'get_out_shape'):
return layer.get_out_shape()
else:
raise Exception('Unhandled base layer type')
def get_out_shape(layer, input_shape):
if hasattr(layer, 'out_channels'):
if type(layer) is torch.nn.modules.conv.Conv1d:
convModel = 1
elif type(layer) is torch.nn.modules.conv.Conv2d:
convModel = 2
return get_output_shape(input_shape,
kernel_size=layer.kernel_size,
stride=layer.stride,
padding=layer.padding,
dilation=layer.dilation,
convModel=convModel)
elif hasattr(layer, 'out_features'):
return []
elif hasattr(layer, 'get_out_shape'):
return layer.get_out_shape()
else:
raise Exception('Unhandled base layer type')