| 113 | |
| 114 | |
| 115 | class PositionalID(torch.nn.Module): |
| 116 | def __init__(self, max_id=25, repeat_length=20): |
| 117 | super().__init__() |
| 118 | self.max_id = max_id |
| 119 | self.repeat_length = repeat_length |
| 120 | |
| 121 | def frame_id_to_position_id(self, frame_id): |
| 122 | if frame_id < self.max_id: |
| 123 | position_id = frame_id |
| 124 | else: |
| 125 | position_id = (frame_id - self.max_id) % (self.repeat_length * 2) |
| 126 | if position_id < self.repeat_length: |
| 127 | position_id = self.max_id - 2 - position_id |
| 128 | else: |
| 129 | position_id = self.max_id - 2 * self.repeat_length + position_id |
| 130 | return position_id |
| 131 | |
| 132 | def forward(self, num_frames, pivot_frame_id=0): |
| 133 | position_ids = [self.frame_id_to_position_id(abs(i-pivot_frame_id)) for i in range(num_frames)] |
| 134 | position_ids = torch.IntTensor(position_ids) |
| 135 | return position_ids |
| 136 | |
| 137 | |
| 138 | class TemporalAttentionBlock(torch.nn.Module): |