| 240 | |
| 241 | class Vec2Patch(nn.Module): |
| 242 | def __init__(self, channel, hidden, output_size, kernel_size, stride, padding): |
| 243 | super(Vec2Patch, self).__init__() |
| 244 | self.relu = nn.LeakyReLU(0.2, inplace=True) |
| 245 | c_out = reduce((lambda x, y: x * y), kernel_size) * channel |
| 246 | self.embedding = nn.Linear(hidden, c_out) |
| 247 | self.to_patch = torch.nn.Fold(output_size=output_size, kernel_size=kernel_size, stride=stride, padding=padding) |
| 248 | h, w = output_size |
| 249 | |
| 250 | def forward(self, x): |
| 251 | feat = self.embedding(x) |