saving of model

ApplyModel.py

@@ -0,0 +1,29 @@
+import numpy as np
+import torch
+from PIL import Image
+
+import DataLoader
+import ex4
+from ImageImpaint import get_train_device
+from netio import load_model
+
+
+def apply_model(filepath: str):
+    device = get_train_device()
+
+    img = Image.open(filepath)
+    model = load_model()
+    model.to(device)
+
+    pic = DataLoader.preprocess(img, precision=np.float32)
+    pic = ex4.ex4(pic, (5, 5), (4, 4))[0]
+    Image.fromarray((np.transpose(pic * 255.0, (1, 2, 0)).astype(np.uint8))).save("filename_grid.jpg")
+    out = model(torch.from_numpy(pic).to(device))
+    out = DataLoader.postprocess(out.cpu().detach().numpy())
+    out = np.transpose(out, (1, 2, 0))
+    im = Image.fromarray(out)
+    im.save("filename.jpg", format="jpeg")
+
+
+if __name__ == '__main__':
+    apply_model("training/000/000017.jpg")

DataLoader.py

@@ -14,22 +14,15 @@ IMG_SIZE = 100
 
 
 class ImageDataset(Dataset):
-    def __init__(self, image_dir):
+    def __init__(self, image_dir, precision: np.float32 or np.float64):
         self.image_files = sorted(glob.glob(os.path.join(image_dir, "**", "*.jpg"), recursive=True))
-        # Mean and std arrays could also be defined as class attributes
+        self.precision = precision
-        # self.norm_mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
-        # self.norm_std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
 
     def __getitem__(self, index):
         # Open image file, convert to numpy array and scale to [0, 1]
         target_image = Image.open(self.image_files[index])
-        # image = np.array(Image.open(self.image_files[index]), dtype=np.float32) / 255
+
-        resize_transforms = transforms.Compose([
+        target_image = preprocess(target_image, self.precision)
-            transforms.Resize(size=IMG_SIZE),
-            transforms.CenterCrop(size=(IMG_SIZE, IMG_SIZE)),
-        ])
-        target_image = resize_transforms(target_image)
-        target_image = preprocess(target_image)
 
         # calculate image with black grid
         doomed_image = ex4.ex4(target_image, (5, 5), (4, 4))
@@ -43,9 +36,16 @@ class ImageDataset(Dataset):
         return len(self.image_files)
 
 
-def preprocess(input: np.array) -> np.array:
+def preprocess(input: np.array, precision: np.float32 or np.float64) -> np.array:
+    # image = np.array(Image.open(self.image_files[index]), dtype=np.float32) / 255
+    resize_transforms = transforms.Compose([
+        transforms.Resize(size=IMG_SIZE),
+        transforms.CenterCrop(size=(IMG_SIZE, IMG_SIZE)),
+    ])
+    input = resize_transforms(input)
+
     # normalize image from 0-1
-    target_image = np.array(input, dtype=np.float64) / 255.0
+    target_image = np.array(input, dtype=precision) / 255.0
 
     # Perform normalization for each channel
     # image = (image - self.norm_mean) / self.norm_std
@@ -59,33 +59,27 @@ def postprocess(input: np.array) -> np.array:
     return target_image
 
 
-def get_image_loader(path: str):
+def get_image_loader(path: str, precision: np.float32 or np.float64):
-    image_dataset = ImageDataset(path)
+    image_dataset = ImageDataset(path, precision)
     totlen = len(image_dataset)
 
-    test_set_size = .001
+    test_set_size = .1
     trains, tests = torch.utils.data.dataset.random_split(image_dataset, lengths=(totlen - int(totlen * test_set_size),
                                                                                   int(totlen * test_set_size)),
-                                                          generator=torch.Generator().manual_seed(42))
+                                                          generator=torch.Generator().manual_seed(0))
 
     train_loader = DataLoader(
         trains,
         shuffle=True,  # shuffle the order of our samples
-        batch_size=5,  # stack 4 samples to a minibatch
+        batch_size=25,  # stack 4 samples to a minibatch
         num_workers=4  # no background workers (see comment below)
     )
 
     test_loader = DataLoader(
         tests,
         shuffle=True,  # shuffle the order of our samples
-        batch_size=1,  # stack 4 samples to a minibatch
+        batch_size=5,  # stack 4 samples to a minibatch
         num_workers=0  # no background workers (see comment below)
     )
 
     return train_loader, test_loader
-
-
-def rgb2gray(rgb_array: np.ndarray):
-    r, g, b = rgb_array[:, :, 0], rgb_array[:, :, 1], rgb_array[:, :, 2]
-    gray = 0.2989 * r + 0.5870 * g + 0.1140 * b
-    return gray

ImageImpaint.py

@@ -1,11 +1,13 @@
 import numpy as np
 import torch
+from PIL.Image import Image
 
+import DataLoader
 from DataLoader import get_image_loader
 from Net import ImageNN
 
 # 01.05.22 -- 0.5h
-from netio import save_model, load_model
+from netio import save_model, load_model, eval_evalset
 
 
 def get_train_device():
@@ -21,17 +23,15 @@ def train_model():
     device = get_train_device()
 
     # Load datasets
-    train_loader, test_loader = get_image_loader("training/")
+    train_loader, test_loader = get_image_loader("training/", precision=np.float32)
-    nn = ImageNN(n_in_channels=3)  # todo pass size ason.
+    nn = ImageNN(n_in_channels=3, precision=np.float32)  # todo net params
     nn.train()  # init with train mode
     nn.to(device)  # send net to device available
 
     optimizer = torch.optim.AdamW(nn.parameters(), lr=0.1, weight_decay=1e-5)  # todo adjust parameters and lr
     loss_function = torch.nn.MSELoss()
-    n_epochs = 10  # todo epcchs here
+    loss_function.to(device)
-
+    n_epochs = 7  # todo epcchs here
-    # todo look wtf is that
-    nn.double()
 
     train_sample_size = len(train_loader)
     losses = []
@@ -40,23 +40,21 @@ def train_model():
         print(f"Epoch {epoch}/{n_epochs}\n")
         i = 0
         for input_tensor, target_tensor in train_loader:
-            input_tensor.to(device)
+            output = nn(input_tensor.to(device))  # get model output (forward pass)
-            target_tensor.to(device)
 
-            output = nn(input_tensor)  # get model output (forward pass)
+            loss = loss_function(output.to(device), target_tensor.to(device))  # compute loss given model output and true target
-            loss = loss_function(output, target_tensor)  # compute loss given model output and true target
             loss.backward()  # compute gradients (backward pass)
             optimizer.step()  # perform gradient descent update step
             optimizer.zero_grad()  # reset gradients
             losses.append(loss.item())
 
+            i += train_loader.batch_size
             print(
                 f'\rTraining epoch {epoch} [{i}/{train_sample_size * train_loader.batch_size}] (curr loss: {loss.item():.3})',
                 end='')
-            i += train_loader.batch_size
 
             # eval model every 3000th sample
-            if i % 15 == 0:
+            if i % 3000 == 0:
                 print(f"\nEvaluating model")
                 eval_loss = eval_model(nn, test_loader, loss_function, device)
                 print(f"Evalution loss={eval_loss}")
@@ -64,8 +62,10 @@ def train_model():
                     best_eval_loss = eval_loss
                     save_model(nn)
 
-    # switch net to eval mode
+                nn.train()
-    print(eval_model(nn, test_loader, loss_function, device=device))
+
+    # evaluate model with submission pkl file
+    eval_evalset()
 
 
 # func to evaluate our trained model
@@ -77,8 +77,8 @@ def eval_model(model: torch.nn.Module, dataloader: torch.utils.data.DataLoader,
     with torch.no_grad():
         i = 0
         for input, target in dataloader:
-            input.to(device)
+            input = input.to(device)
-            target.to(device)
+            target = target.to(device)
 
             out = model(input)
             loss += loss_fn(out, target).item()
@@ -86,11 +86,9 @@ def eval_model(model: torch.nn.Module, dataloader: torch.utils.data.DataLoader,
             i += dataloader.batch_size
     print()
     loss /= len(dataloader)
-    model.train()
     return loss
 
 
-def apply_model():
-    model = load_model()
 
-    pass
+if __name__ == '__main__':
+    train_model()

Net.py

@@ -1,8 +1,9 @@
+import numpy as np
 import torch
 
 
 class ImageNN(torch.nn.Module):
-    def __init__(self, n_in_channels: int = 1, n_hidden_layers: int = 3, n_kernels: int = 32, kernel_size: int = 7):
+    def __init__(self, precision: np.float32 or np.float64, n_in_channels: int = 1, n_hidden_layers: int = 3, n_kernels: int = 32, kernel_size: int = 7):
         """Simple CNN with `n_hidden_layers`, `n_kernels`, and `kernel_size` as hyperparameters"""
         super().__init__()
 
@@ -25,6 +26,9 @@ class ImageNN(torch.nn.Module):
             padding=int(kernel_size / 2)
         )
 
+        if precision == np.float64:
+            self.double()
+
     def forward(self, x):
         """Apply CNN to input `x` of shape (N, n_channels, X, Y), where N=n_samples and X, Y are spatial dimensions"""
         cnn_out = self.hidden_layers(x)  # apply hidden layers (N, n_in_channels, X, Y) -> (N, n_kernels, X, Y)

main.py

@@ -1,11 +0,0 @@
-# This is a sample Python script.
-
-# Press Umschalt+F10 to execute it or replace it with your code.
-# Press Double Shift to search everywhere for classes, files, tool windows, actions, and settings.
-
-
-# See PyCharm help at https://www.jetbrains.com/help/pycharm/
-from ImageImpaint import train_model
-
-if __name__ == '__main__':
-    train_model()

netio.py

@@ -1,5 +1,5 @@
+import os
 import pickle
-import sys
 
 import numpy as np
 import torch
@@ -15,32 +15,40 @@ def save_model(model: torch.nn.Module):
     print(f"Saved raw model to {MODEL_PATH}")
     torch.save(model, MODEL_PATH)
 
+
+def eval_evalset():
     # read the provided testing pickle file
     print("Generating pickle file with privided test data")
+    try:
+        os.unlink(PICKEL_PATH)
+    except:
+        pass
+
+    model = load_model()
     model.eval()
     with open('testing/inputs.pkl', 'rb') as handle:
-        with open(PICKEL_PATH, 'wb') as writehandle:
         b: dict = pickle.load(handle)
-            outarr = []
+        outarr = np.zeros(dtype=np.uint8, shape=(len(b['input_arrays']), 3, 100, 100))
         i = 0
         piclen = len(b['input_arrays'])
         for pic in b['input_arrays']:
-                pic = DataLoader.preprocess(pic)
+            pic = DataLoader.preprocess(pic, precision=np.float32)
             out = model(torch.from_numpy(pic))
-                out = DataLoader.postprocess(out.detach().numpy())
+            out = DataLoader.postprocess(out.cpu().detach().numpy())
-                pickle.dump(out, writehandle, protocol=pickle.HIGHEST_PROTOCOL)
+            outarr[i] = out
 
-                print(
+            print(f'\rApplying model [{i}/{piclen}]', end='')
-                    f'\rApplying model [{i}/{piclen}] {sys.getsizeof(outarr)}',end='')
             i += 1
 
+        write_to_pickle(PICKEL_PATH, list(outarr))
     # compress the generated pickle arr
     Compress.compress(PICKEL_PATH)
 
 
-def load_model():
+def write_to_pickle(filename: str, data):
-    model = ImageNN()
+    with open(filename, 'wb') as handle:
-    model.load_state_dict(torch.load(MODEL_PATH))
+        pickle.dump(data, handle, protocol=pickle.HIGHEST_PROTOCOL)
-    model.eval()
-    return model
 
+
+def load_model():
+    return torch.load(MODEL_PATH)