| 383 | |
| 384 | |
| 385 | class AffineGridGen(Module): |
| 386 | def __init__(self, out_h=240, out_w=240, out_ch=3, use_cuda=True): |
| 387 | super(AffineGridGen, self).__init__() |
| 388 | self.out_h = out_h |
| 389 | self.out_w = out_w |
| 390 | self.out_ch = out_ch |
| 391 | |
| 392 | def forward(self, theta): |
| 393 | b = theta.size()[0] |
| 394 | if not theta.size() == (b, 2, 3): |
| 395 | theta = theta.view(-1, 2, 3) |
| 396 | theta = theta.contiguous() |
| 397 | batch_size = theta.size()[0] |
| 398 | out_size = torch.Size((batch_size, self.out_ch, self.out_h, self.out_w)) |
| 399 | return F.affine_grid(theta, out_size, align_corners=True) |
| 400 | |
| 401 | |
| 402 | class AffineGridGenV2(Module): |