| 4 | |
| 5 | |
| 6 | class classifier(nn.Module): |
| 7 | def __init__(self, ef_dim=32, z_dim=512, class_num=24, voxel_size=128): |
| 8 | super(classifier, self).__init__() |
| 9 | self.ef_dim = ef_dim |
| 10 | self.z_dim = z_dim |
| 11 | self.class_num = class_num |
| 12 | self.voxel_size = voxel_size |
| 13 | |
| 14 | self.conv_1 = nn.Conv3d(1, self.ef_dim, 4, stride=2, padding=1, bias=True) |
| 15 | self.bn_1 = nn.InstanceNorm3d(self.ef_dim) |
| 16 | |
| 17 | self.conv_2 = nn.Conv3d(self.ef_dim, self.ef_dim*2, 4, stride=2, padding=1, bias=True) |
| 18 | self.bn_2 = nn.InstanceNorm3d(self.ef_dim*2) |
| 19 | |
| 20 | self.conv_3 = nn.Conv3d(self.ef_dim*2, self.ef_dim*4, 4, stride=2, padding=1, bias=True) |
| 21 | self.bn_3 = nn.InstanceNorm3d(self.ef_dim*4) |
| 22 | |
| 23 | self.conv_4 = nn.Conv3d(self.ef_dim*4, self.ef_dim*8, 4, stride=2, padding=1, bias=True) |
| 24 | self.bn_4 = nn.InstanceNorm3d(self.ef_dim*8) |
| 25 | |
| 26 | self.conv_5 = nn.Conv3d(self.ef_dim*8, self.z_dim, 4, stride=2, padding=1, bias=True) |
| 27 | |
| 28 | if self.voxel_size==256: |
| 29 | self.bn_5 = nn.InstanceNorm3d(self.z_dim) |
| 30 | self.conv_5_2 = nn.Conv3d(self.z_dim, self.z_dim, 4, stride=2, padding=1, bias=True) |
| 31 | |
| 32 | self.linear1 = nn.Linear(self.z_dim, self.class_num, bias=True) |
| 33 | |
| 34 | def forward(self, inputs, out_layer=None, is_training=False): |
| 35 | out = inputs |
| 36 | |
| 37 | out = self.bn_1(self.conv_1(out)) |
| 38 | out = F.leaky_relu(out, negative_slope=0.01, inplace=True) |
| 39 | |
| 40 | if out_layer == 1: |
| 41 | return out |
| 42 | |
| 43 | out = self.bn_2(self.conv_2(out)) |
| 44 | out = F.leaky_relu(out, negative_slope=0.01, inplace=True) |
| 45 | |
| 46 | if out_layer == 2: |
| 47 | return out |
| 48 | |
| 49 | out = self.bn_3(self.conv_3(out)) |
| 50 | out = F.leaky_relu(out, negative_slope=0.01, inplace=True) |
| 51 | |
| 52 | if out_layer == 3: |
| 53 | return out |
| 54 | |
| 55 | out = self.bn_4(self.conv_4(out)) |
| 56 | out = F.leaky_relu(out, negative_slope=0.01, inplace=True) |
| 57 | |
| 58 | if out_layer == 4: |
| 59 | return out |
| 60 | |
| 61 | if self.voxel_size==256: |
| 62 | out = self.bn_5(out) |
| 63 | out = F.leaky_relu(out, negative_slope=0.01, inplace=True) |
no outgoing calls
no test coverage detected