12345678910111213141516171819202122232425262728293031323334 |
- from __future__ import absolute_import
- from __future__ import division
- from __future__ import print_function
- from paddle import nn
- from paddle.nn import functional as F
- class PRENHead(nn.Layer):
- def __init__(self, in_channels, out_channels, **kwargs):
- super(PRENHead, self).__init__()
- self.linear = nn.Linear(in_channels, out_channels)
- def forward(self, x, targets=None):
- predicts = self.linear(x)
- if not self.training:
- predicts = F.softmax(predicts, axis=2)
- return predicts
|