Method forward

model/linear_probe.py:215–252

(self, x, sensor_type = None, return_feature = False)

        trunc_normal_(self.head.weight, std=0.01)
    
    def forward(self, x, sensor_type = None, return_feature = False):

        with torch.no_grad():
            if self.pooling == 'none':
                x = self.tactile_model(x, sensor_type = sensor_type, probe=True)
                out = x
            else:
                if self.pooling == 'cls':
                    # out = self.touch_projection(x.pooler_output)
                    x = self.tactile_model(x, sensor_type = sensor_type, probe=True, get_cls=True)
                    out = x
                elif self.pooling == 'last':
                    x = self.tactile_model(x, sensor_type = sensor_type, probe=True, get_cls=False)
                    out = x[:, -self.single_patch_num:, :]
                else:
                    x = self.tactile_model(x, sensor_type = sensor_type, probe=True, get_cls=False)
                    if self.use_sensor_token:
                        out = x[:, 6:, :]
                    else:
                        out = x[:, 1:, :]

        feature = out
        # print(out.shape)

        if self.pooling == 'none':
            out = self.head(out)

        else:
            if self.pooling == 'cls':
                out = self.head(out)

            elif self.pooling == 'global' or self.pooling == 'last':
                out = self.head(out.mean(dim=1))

        
        if return_feature:
            return out, feature
        return out