vehicle_neck.py 1.6 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152
  1. # Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. from __future__ import absolute_import, division, print_function
  15. import paddle
  16. import paddle.nn as nn
  17. class VehicleNeck(nn.Layer):
  18. def __init__(self,
  19. in_channels,
  20. out_channels,
  21. kernel_size=1,
  22. stride=1,
  23. padding=0,
  24. dilation=1,
  25. groups=1,
  26. padding_mode='zeros',
  27. weight_attr=None,
  28. bias_attr=None,
  29. data_format='NCHW'):
  30. super().__init__()
  31. self.conv = nn.Conv2D(
  32. in_channels=in_channels,
  33. out_channels=out_channels,
  34. kernel_size=kernel_size,
  35. stride=stride,
  36. padding=padding,
  37. dilation=dilation,
  38. groups=groups,
  39. padding_mode=padding_mode,
  40. weight_attr=weight_attr,
  41. bias_attr=weight_attr,
  42. data_format=data_format)
  43. self.flatten = nn.Flatten()
  44. def forward(self, x):
  45. x = self.conv(x)
  46. x = self.flatten(x)
  47. return x