(self, args, config, num_frames, tube_size)
| 151 | |
| 152 | class TactileProbeVideo(nn.Module): |
| 153 | def __init__(self, args, config, num_frames, tube_size): |
| 154 | super(TactileProbeVideo, self).__init__() |
| 155 | |
| 156 | config.vision_config.num_frames = num_frames |
| 157 | config.vision_config.tube_size = tube_size # Unused |
| 158 | |
| 159 | self.tactile_model = TactileVideoMAE(args, config, num_frames, tube_size) |
| 160 | |
| 161 | self.pooling = args.pooling |
| 162 | |
| 163 | if args.dataset == 'material': |
| 164 | self.classes = 20 |
| 165 | elif args.dataset == 'cloth': |
| 166 | self.classes = 20 |
| 167 | elif args.dataset == 'rough': |
| 168 | self.classes = 1 |
| 169 | elif args.dataset == 'hard': |
| 170 | self.classes = 1 |
| 171 | elif args.dataset == 'obj2' or args.dataset == 'obj1' or args.dataset == 'objreal': |
| 172 | self.classes = 7 |
| 173 | else: |
| 174 | print('sparsh bench') |
| 175 | |
| 176 | if 'force' in args.dataset: |
| 177 | self.head = ForceHead( |
| 178 | embed_dim=config.vision_config.hidden_size, |
| 179 | num_heads=16, |
| 180 | mlp_ratio=4.0, |
| 181 | depth=1, |
| 182 | norm_layer=nn.LayerNorm, |
| 183 | init_std=0.02, |
| 184 | qkv_bias=True, |
| 185 | complete_block=True |
| 186 | ) |
| 187 | self.pooling = 'none' |
| 188 | elif args.dataset in ['cloth']: |
| 189 | self.head = AttnHead( |
| 190 | embed_dim=config.vision_config.hidden_size, |
| 191 | num_heads=16, |
| 192 | mlp_ratio=4.0, |
| 193 | depth=1, |
| 194 | norm_layer=nn.LayerNorm, |
| 195 | init_std=0.02, |
| 196 | qkv_bias=True, |
| 197 | complete_block=True, |
| 198 | classes = self.classes |
| 199 | ) |
| 200 | self.pooling = 'none' |
| 201 | else: |
| 202 | if self.pooling == 'cls': |
| 203 | self.head = nn.Linear(config.projection_dim, self.classes) |
| 204 | else: |
| 205 | self.head = nn.Linear(config.vision_config.hidden_size, self.classes) |
| 206 | |
| 207 | self.dataset = args.dataset |
| 208 | |
| 209 | self.single_patch_num = (config.vision_config.image_size // config.vision_config.patch_size) ** 2 |
| 210 |
nothing calls this directly
no test coverage detected