21 lines
600 B
Python
21 lines
600 B
Python
from torch import nn
|
|
import PytorchBoot.stereotype as stereotype
|
|
|
|
@stereotype.module("pts_num_encoder")
|
|
class PointsNumEncoder(nn.Module):
|
|
def __init__(self, config):
|
|
super(PointsNumEncoder, self).__init__()
|
|
self.config = config
|
|
out_dim = config["out_dim"]
|
|
self.act = nn.ReLU(True)
|
|
|
|
self.pts_num_encoder = nn.Sequential(
|
|
nn.Linear(1, out_dim),
|
|
self.act,
|
|
nn.Linear(out_dim, out_dim),
|
|
self.act,
|
|
)
|
|
|
|
def encode_pts_num(self, num_seq):
|
|
return self.pts_num_encoder(num_seq)
|