| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from model.common import (MLP, Custom1x1Subm3d, ResidualBlock, UBlock) |
|
|
| import functools |
| from collections import OrderedDict |
| import spconv.pytorch as spconv |
| from spconv.pytorch.modules import SparseModule |
|
|
| |
| __all__ = ['Encoder', 'Decoder',] |
| |
| class Encoder(nn.Module): |
| def __init__(self, input_dim, hidden_dim, spconv_channels, num_blocks): |
| super(Encoder, self).__init__() |
| |
| block = ResidualBlock |
| norm_fn = functools.partial(nn.BatchNorm1d, eps=1e-4, momentum=0.1) |
|
|
| self.input_conv = spconv.SparseSequential( |
| spconv.SubMConv3d( |
| input_dim, spconv_channels, kernel_size=3, padding=1, bias=False, indice_key='subm1')) |
| block_channels = [spconv_channels * (i + 1) for i in range(num_blocks)] |
| self.unet = UBlock(block_channels, norm_fn, 2, block, indice_key_id=1) |
| self.output_layer = spconv.SparseSequential(norm_fn(spconv_channels), nn.ReLU()) |
|
|
|
|
| def forward(self, voxel_feats, voxel_coords, spatial_shape): |
|
|
| |
| batch_size = len(voxel_coords) |
| voxel_feats_total = [] |
| voxel_coords_total = [] |
| voxel_batch_id_total = [] |
| for i in range(batch_size): |
| batch_col = torch.zeros((voxel_coords[i].shape[0], 1), device=voxel_coords[i].device, dtype=torch.int32) |
| voxel_batch_id_total.append(batch_col) |
| voxel_coords_total.append(voxel_coords[i].int()) |
| voxel_coords[i] = torch.cat([batch_col, voxel_coords[i].int()], dim=1) |
| input = spconv.SparseConvTensor(voxel_feats[i], voxel_coords[i], spatial_shape, 1) |
| output = self.input_conv(input) |
| output = self.unet(output) |
| voxel_feats_total.append(output) |
|
|
| return voxel_feats_total, voxel_coords_total, voxel_batch_id_total |
|
|
| class Decoder(nn.Module): |
| def __init__(self, input_dim, hidden_dim, output_dim, num_layers): |
| super(Decoder, self).__init__() |
| self.num_layers = num_layers |
| self.start_layer = nn.Sequential( |
| nn.Linear(input_dim, hidden_dim), |
| nn.ReLU(), |
| nn.Linear(hidden_dim, hidden_dim) |
| ) |
| self.layers = nn.ModuleList([ |
| nn.Sequential( |
| nn.Linear(hidden_dim, hidden_dim), |
| nn.ReLU(), |
| nn.Linear(hidden_dim, hidden_dim) |
| ) for _ in range(num_layers - 2) |
| ]) |
| self.final_layer = nn.Sequential( |
| nn.Linear(hidden_dim, hidden_dim), |
| nn.ReLU(), |
| nn.Linear(hidden_dim, output_dim) |
| ) |
|
|
| def forward(self, x): |
| x = self.start_layer(x) |
| for layer in self.layers: |
| x = layer(x) |
| return self.final_layer(x) |
