(
self,
embed_dim=256,
kernel_size=4,
stride=4,
padding=0,
total_stride=16,
activation=nn.GELU,
)
| 24 | """ |
| 25 | |
| 26 | def __init__( |
| 27 | self, |
| 28 | embed_dim=256, |
| 29 | kernel_size=4, |
| 30 | stride=4, |
| 31 | padding=0, |
| 32 | total_stride=16, |
| 33 | activation=nn.GELU, |
| 34 | ): |
| 35 | super().__init__() |
| 36 | num_layers = int(math.log2(total_stride) // math.log2(stride)) |
| 37 | assert stride**num_layers == total_stride |
| 38 | self.encoder = nn.Sequential() |
| 39 | mask_in_chans, mask_out_chans = 1, 1 |
| 40 | for _ in range(num_layers): |
| 41 | mask_out_chans = mask_in_chans * (stride**2) |
| 42 | self.encoder.append( |
| 43 | nn.Conv2d( |
| 44 | mask_in_chans, |
| 45 | mask_out_chans, |
| 46 | kernel_size=kernel_size, |
| 47 | stride=stride, |
| 48 | padding=padding, |
| 49 | ) |
| 50 | ) |
| 51 | self.encoder.append(LayerNorm2d(mask_out_chans)) |
| 52 | self.encoder.append(activation()) |
| 53 | mask_in_chans = mask_out_chans |
| 54 | |
| 55 | self.encoder.append(nn.Conv2d(mask_out_chans, embed_dim, kernel_size=1)) |
| 56 | |
| 57 | def forward(self, x): |
| 58 | return self.encoder(x) |
nothing calls this directly
no test coverage detected