| 172 | return embeddings |
| 173 | |
| 174 | class _EncoderBlock(nn.Module): |
| 175 | def __init__(self, in_channels, out_channels, downsample=True): |
| 176 | super(_EncoderBlock, self).__init__() |
| 177 | self.downsample = downsample |
| 178 | self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) |
| 179 | layers = [ |
| 180 | conv3x3(in_channels, out_channels), |
| 181 | nn.BatchNorm2d(out_channels), |
| 182 | nn.ReLU(inplace=True), |
| 183 | conv3x3(out_channels, out_channels), |
| 184 | nn.BatchNorm2d(out_channels), |
| 185 | nn.ReLU(inplace=True), |
| 186 | ] |
| 187 | self.encode = nn.Sequential(*layers) |
| 188 | |
| 189 | def forward(self, x): |
| 190 | if self.downsample: |
| 191 | x = self.maxpool(x) |
| 192 | x = self.encode(x) |
| 193 | return x |
| 194 | |
| 195 | class Transformer(nn.Module): |
| 196 | def __init__(self, in_channels1, in_channels2, feat_size1, feat_size2, hidden_size=args['D']): |