text: # Copyright 2020 - 2021 MONAI Consortium
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #     http://www.apache.org/licenses/LICENSE-2.0
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  
  import unittest
  
  import numpy as np
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  import torch.optim as optim
  
  from monai.data import create_test_image_2d
  from monai.networks.layers import AffineTransform
  from monai.utils import set_determinism
  from tests.utils import DistTestCase, TimedCall
  
  
  class STNBenchmark(nn.Module):
      """
      adapted from https://pytorch.org/tutorials/intermediate/spatial_transformer_tutorial.html
      """
  
      def __init__(self, is_ref=True, reverse_indexing=False):
          super().__init__()
          self.is_ref = is_ref
          self.localization = nn.Sequential(
              nn.Conv2d(1, 8, kernel_size=7),
              nn.MaxPool2d(2, stride=2),
              nn.ReLU(True),
              nn.Conv2d(8, 10, kernel_size=5),
              nn.MaxPool2d(2, stride=2),
              nn.ReLU(True),
          )
          # Regressor for the 3 * 2 affine matrix
          self.fc_loc = nn.Sequential(nn.Linear(10 * 3 * 3, 32), nn.ReLU(True), nn.Linear(32, 3 * 2))
          # Initialize the weights/bias with identity transformation
          self.fc_loc[2].weight.data.zero_()
          self.fc_loc[2].bias.data.copy_(torch.tensor([1, 0, 0, 0, 1, 0], dtype=torch.float))
          if not self.is_ref:
              self.xform = AffineTransform(normalized=True, reverse_indexing=reverse_indexing)
  
      # Spatial transformer network forward function
      def stn_ref(self, x):
          xs = self.localization(x)
          xs = xs.view(-1, 10 * 3 * 3)
          theta = self.fc_loc(xs)
          theta = theta.view(-1, 2, 3)
  
          grid = F.affine_grid(theta, x.size(), align_corners=False)
          x = F.grid_sample(x, grid, align_corners=False)
          return x
  
      def stn(self, x):
          xs = self.localization(x)
          xs = xs.view(-1, 10 * 3 * 3)
          theta = self.fc_loc(xs)
          theta = theta.view(-1, 2, 3)
          x = self.xform(x, theta, spatial_size=x.size()[2:])
          return x
  
      def forward(self, x):
          if self.is_ref:
              return self.stn_ref(x)
          return self.stn(x)
  
  
  def compare_2d(is_ref=True, device=None, reverse_indexing=False):
      batch_size = 32
      img_a = [create_test_image_2d(28, 28, 5, rad_max=6, noise_max=1)[0][None] for _ in range(batch_size)]
      img_b = [create_test_image_2d(28, 28, 5, rad_max=6, noise_max=1)[0][None] for _ in range(batch_size)]
      img_a = np.stack(img_a, axis=0)
      img_b = np.stack(img_b, axis=0)
      img_a = torch.as_tensor(img_a, device=device)
      img_b = torch.as_tensor(img_b, device=device)
      model = STNBenchmark(is_ref=is_ref, reverse_indexing=reverse_indexing).to(device)
      optimizer = optim.SGD(model.parameters(), lr=0.001)
      model.train()
      init_loss = None
      for _ in range(20):
          optimizer.zero_grad()
          output_a = model(img_a)
          loss = torch.mean((output_a - img_b) ** 2)
          if init_loss is None:
              init_loss = loss.item()
          loss.backward()
          optimizer.step()
      return model(img_a).detach().cpu().numpy(), loss.item(), init_loss
  
  
  class TestSpatialTransformerCore(DistTestCase):
      def setUp(self):
          self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu:0")
  
      def tearDown(self):
          set_determinism(seed=None)
  
      @TimedCall(seconds=60)
      def test_training(self):
          """
          check that the quality AffineTransform backpropagation
          """
          atol = 1e-5
          set_determinism(seed=0)
          out_ref, loss_ref, init_loss_ref = compare_2d(True, self.device)
          print(out_ref.shape, loss_ref, init_loss_ref)
  
          set_determinism(seed=0)
          out, loss, init_loss = compare_2d(False, self.device)
          print(out.shape, loss, init_loss)
          np.testing.assert_allclose(out_ref, out, atol=atol)
          np.testing.assert_allclose(init_loss_ref, init_loss, atol=atol)
          np.testing.assert_allclose(loss_ref, loss, atol=atol)
  
          set_determinism(seed=0)
          out, loss, init_loss = compare_2d(False, self.device, True)
          print(out.shape, loss, init_loss)
          np.testing.assert_allclose(out_ref, out, atol=atol)
          np.testing.assert_allclose(init_loss_ref, init_loss, atol=atol)
          np.testing.assert_allclose(loss_ref, loss, atol=atol)
  
  
  if __name__ == "__main__":
      unittest.main()
  