| 159 | return feature_global.reshape(bs, g, self.encoder_channel) |
| 160 | |
| 161 | class PointcloudEncoder(nn.Module): |
| 162 | def __init__(self, point_transformer, args): |
| 163 | super().__init__() |
| 164 | from easydict import EasyDict |
| 165 | self.trans_dim = args.pc_feat_dim # 768 |
| 166 | self.embed_dim = args.embed_dim # 512 |
| 167 | self.group_size = args.group_size # 32 |
| 168 | self.num_group = args.num_group # 512 |
| 169 | # grouper |
| 170 | self.group_divider = Group(num_group = self.num_group, group_size = self.group_size) |
| 171 | # define the encoder |
| 172 | self.encoder_dim = args.pc_encoder_dim # 256 |
| 173 | self.encoder = Encoder(encoder_channel = self.encoder_dim) |
| 174 | |
| 175 | # bridge encoder and transformer |
| 176 | self.encoder2trans = nn.Linear(self.encoder_dim, self.trans_dim) |
| 177 | |
| 178 | # bridge transformer and clip embedding |
| 179 | self.trans2embed = nn.Linear(self.trans_dim, self.embed_dim) |
| 180 | self.cls_token = nn.Parameter(torch.zeros(1, 1, self.trans_dim)) |
| 181 | self.cls_pos = nn.Parameter(torch.randn(1, 1, self.trans_dim)) |
| 182 | |
| 183 | self.pos_embed = nn.Sequential( |
| 184 | nn.Linear(3, 128), |
| 185 | nn.GELU(), |
| 186 | nn.Linear(128, self.trans_dim) |
| 187 | ) |
| 188 | # setting a patch_dropout of 0. would mean it is disabled and this function would be the identity fn |
| 189 | self.patch_dropout = PatchDropout(args.patch_dropout) if args.patch_dropout > 0. else nn.Identity() |
| 190 | self.visual = point_transformer |
| 191 | |
| 192 | |
| 193 | def forward(self, pts, colors): |
| 194 | # divide the point cloud in the same form. This is important |
| 195 | _, center, features = self.group_divider(pts, colors) |
| 196 | |
| 197 | # encoder the input cloud patches |
| 198 | group_input_tokens = self.encoder(features) # B G N |
| 199 | group_input_tokens = self.encoder2trans(group_input_tokens) |
| 200 | # prepare cls |
| 201 | cls_tokens = self.cls_token.expand(group_input_tokens.size(0), -1, -1) |
| 202 | cls_pos = self.cls_pos.expand(group_input_tokens.size(0), -1, -1) |
| 203 | # add pos embedding |
| 204 | pos = self.pos_embed(center) |
| 205 | # final input |
| 206 | x = torch.cat((cls_tokens, group_input_tokens), dim=1) |
| 207 | pos = torch.cat((cls_pos, pos), dim=1) |
| 208 | # transformer |
| 209 | x = x + pos |
| 210 | # x = x.half() |
| 211 | |
| 212 | # a patch_dropout of 0. would mean it is disabled and this function would do nothing but return what was passed in |
| 213 | x = self.patch_dropout(x) |
| 214 | |
| 215 | x = self.visual.pos_drop(x) |
| 216 | |
| 217 | # ModuleList not support forward |
| 218 | for i, blk in enumerate(self.visual.blocks): |