| 170 | |
| 171 | class Tokenizer(nn.Module): |
| 172 | def __init__(self, |
| 173 | kernel_size, stride, padding, |
| 174 | pooling_kernel_size=3, pooling_stride=2, pooling_padding=1, |
| 175 | n_conv_layers=1, |
| 176 | n_input_channels=3, |
| 177 | n_output_channels=64, |
| 178 | in_planes=64, |
| 179 | activation=None, |
| 180 | max_pool=True, |
| 181 | conv_bias=False): |
| 182 | super().__init__() |
| 183 | |
| 184 | n_filter_list = [n_input_channels] + \ |
| 185 | [in_planes for _ in range(n_conv_layers - 1)] + \ |
| 186 | [n_output_channels] |
| 187 | |
| 188 | n_filter_list_pairs = zip(n_filter_list[:-1], n_filter_list[1:]) |
| 189 | |
| 190 | self.conv_layers = nn.Sequential( |
| 191 | *[nn.Sequential( |
| 192 | nn.Conv2d(chan_in, chan_out, |
| 193 | kernel_size=(kernel_size, kernel_size), |
| 194 | stride=(stride, stride), |
| 195 | padding=(padding, padding), bias=conv_bias), |
| 196 | nn.Identity() if not exists(activation) else activation(), |
| 197 | nn.MaxPool2d(kernel_size=pooling_kernel_size, |
| 198 | stride=pooling_stride, |
| 199 | padding=pooling_padding) if max_pool else nn.Identity() |
| 200 | ) |
| 201 | for chan_in, chan_out in n_filter_list_pairs |
| 202 | ]) |
| 203 | |
| 204 | self.apply(self.init_weight) |
| 205 | |
| 206 | def sequence_length(self, n_channels=3, height=224, width=224): |
| 207 | return self.forward(torch.zeros((1, n_channels, height, width))).shape[1] |