adapted from https://pytorch.org/tutorials/intermediate/spatial_transformer_tutorial.html
| 26 | |
| 27 | |
| 28 | class STNBenchmark(nn.Module): |
| 29 | """ |
| 30 | adapted from https://pytorch.org/tutorials/intermediate/spatial_transformer_tutorial.html |
| 31 | """ |
| 32 | |
| 33 | def __init__(self, is_ref=True, reverse_indexing=False): |
| 34 | super().__init__() |
| 35 | self.is_ref = is_ref |
| 36 | self.localization = nn.Sequential( |
| 37 | nn.Conv2d(1, 8, kernel_size=7), |
| 38 | nn.MaxPool2d(2, stride=2), |
| 39 | nn.ReLU(True), |
| 40 | nn.Conv2d(8, 10, kernel_size=5), |
| 41 | nn.MaxPool2d(2, stride=2), |
| 42 | nn.ReLU(True), |
| 43 | ) |
| 44 | # Regressor for the 3 * 2 affine matrix |
| 45 | self.fc_loc = nn.Sequential(nn.Linear(10 * 3 * 3, 32), nn.ReLU(True), nn.Linear(32, 3 * 2)) |
| 46 | # Initialize the weights/bias with identity transformation |
| 47 | self.fc_loc[2].weight.data.zero_() |
| 48 | self.fc_loc[2].bias.data.copy_(torch.tensor([1, 0, 0, 0, 1, 0], dtype=torch.float)) |
| 49 | if not self.is_ref: |
| 50 | self.xform = AffineTransform(align_corners=False, normalized=True, reverse_indexing=reverse_indexing) |
| 51 | |
| 52 | # Spatial transformer network forward function |
| 53 | def stn_ref(self, x): |
| 54 | xs = self.localization(x) |
| 55 | xs = xs.view(-1, 10 * 3 * 3) |
| 56 | theta = self.fc_loc(xs) |
| 57 | theta = theta.view(-1, 2, 3) |
| 58 | |
| 59 | grid = F.affine_grid(theta, x.size(), align_corners=False) |
| 60 | x = F.grid_sample(x, grid, align_corners=False) |
| 61 | return x |
| 62 | |
| 63 | def stn(self, x): |
| 64 | xs = self.localization(x) |
| 65 | xs = xs.view(-1, 10 * 3 * 3) |
| 66 | theta = self.fc_loc(xs) |
| 67 | theta = theta.view(-1, 2, 3) |
| 68 | x = self.xform(x, theta, spatial_size=x.size()[2:]) |
| 69 | return x |
| 70 | |
| 71 | def forward(self, x): |
| 72 | if self.is_ref: |
| 73 | return self.stn_ref(x) |
| 74 | return self.stn(x) |
| 75 | |
| 76 | |
| 77 | def compare_2d(is_ref=True, device=None, reverse_indexing=False): |
no outgoing calls
searching dependent graphs…