| 9 | class TransNetV2(nn.Module): |
| 10 | |
| 11 | def __init__(self, |
| 12 | F=16, L=3, S=2, D=1024, |
| 13 | use_many_hot_targets=True, |
| 14 | use_frame_similarity=True, |
| 15 | use_color_histograms=True, |
| 16 | use_mean_pooling=False, |
| 17 | dropout_rate=0.5, |
| 18 | use_convex_comb_reg=False, # not supported |
| 19 | use_resnet_features=False, # not supported |
| 20 | use_resnet_like_top=False, # not supported |
| 21 | frame_similarity_on_last_layer=False): # not supported |
| 22 | super(TransNetV2, self).__init__() |
| 23 | |
| 24 | if use_resnet_features or use_resnet_like_top or use_convex_comb_reg or frame_similarity_on_last_layer: |
| 25 | raise NotImplemented( |
| 26 | "Some options not implemented in Pytorch version of Transnet!") |
| 27 | |
| 28 | self.SDDCNN = nn.ModuleList( |
| 29 | [StackedDDCNNV2(in_filters=3, n_blocks=S, filters=F, stochastic_depth_drop_prob=0.)] + |
| 30 | [StackedDDCNNV2(in_filters=(F * 2 ** (i - 1)) * 4, |
| 31 | n_blocks=S, filters=F * 2 ** i) for i in range(1, L)] |
| 32 | ) |
| 33 | |
| 34 | self.frame_sim_layer = FrameSimilarity( |
| 35 | sum([(F * 2 ** i) * 4 for i in range(L)]), lookup_window=101, output_dim=128, similarity_dim=128, use_bias=True |
| 36 | ) if use_frame_similarity else None |
| 37 | self.color_hist_layer = ColorHistograms( |
| 38 | lookup_window=101, output_dim=128 |
| 39 | ) if use_color_histograms else None |
| 40 | |
| 41 | self.dropout = nn.Dropout( |
| 42 | dropout_rate) if dropout_rate is not None else None |
| 43 | |
| 44 | output_dim = ((F * 2 ** (L - 1)) * 4) * 3 * \ |
| 45 | 6 # 3x6 for spatial dimensions |
| 46 | if use_frame_similarity: |
| 47 | output_dim += 128 |
| 48 | if use_color_histograms: |
| 49 | output_dim += 128 |
| 50 | |
| 51 | self.fc1 = nn.Linear(output_dim, D) |
| 52 | self.cls_layer1 = nn.Linear(D, 1) |
| 53 | self.cls_layer2 = nn.Linear(D, 1) if use_many_hot_targets else None |
| 54 | |
| 55 | self.use_mean_pooling = use_mean_pooling |
| 56 | self.eval() |
| 57 | |
| 58 | def forward(self, inputs): |
| 59 | assert isinstance(inputs, torch.Tensor) and list(inputs.shape[2:]) == [27, 48, 3] and inputs.dtype == torch.uint8, \ |