Correct image normalization

.gitignore

@@ -8,3 +8,4 @@ __pycache__/
 *.out
 *.pdf
 .DS_Store
+hw2/code/checkpoints/

hw2/code/datasets.py

@@ -44,7 +44,7 @@ def get_data_loader(
         transforms.Resize(image_size),
         transforms.ToImage(),
         transforms.ToDtype(torch.float32, scale=True),
-        transforms.Normalize(mean=[-127.0, -127.0, -127.0], std=[128.0, 128.0, 128.0])
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
     ]
 
     # You should insert some data augmentation techniques to `data_transforms` when `augment` is True
@@ -58,8 +58,6 @@ def get_data_loader(
     # Use `transforms.Compose` to compose the list of transforms into a single transform
     data_transforms = transforms.Compose(data_transforms)
 
-    print(type(data_transforms))
-
     # >>> TODO 1.2: Define the dataset.
     # You should build the path to the selected dataset according to the `mode` parameter,
     # and use the `ImageFolder` class from `torchvision.datasets` to load the datasets.

hw2/code/networks.py

@@ -187,7 +187,7 @@ class Classifier(nn.Module):
 
         # Step 3: use `Tensor.view()` to flatten the tensor to match the size of the input of the
         # fully connected layers.
-        x = x.view(-1, 2048)
+        x = x.view(x.shape[0], -1)
 
         # Step 4: forward process for the fully connected network
         out = self.fc_net(x)
@@ -241,8 +241,8 @@ class STN(nn.Module):
         # Suggested structure: 2 linear layers with one BN and ReLU.
         self.localization_fc = nn.Sequential(
             nn.Linear(16, 256),
-            nn.Linear(256, 360),
-            nn.BatchNorm1d(360),
+            nn.Linear(256, 6),
+            nn.BatchNorm1d(6),
             nn.ReLU()
         )
         # <<< TODO 4.1