Add train_dlaa.py for model training

train_dlaa.py

@@ -0,0 +1,112 @@
+# Created by MSXYZ
+
+### How to train your DLAA model?
+# I have included the `train_dlaa.py` script for those who want to fine-tune the model. 
+# 1. Create a folder named `dataset` next to the script.
+# 2. Put 50-100 high-quality images (landscapes, portraits, geometric images) inside it.
+# 3. Open the terminal in hte folder Run "python train_dlaa.py" and wait for 100 epochs.
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torch.utils.data import DataLoader, Dataset
+from torchvision import transforms
+from PIL import Image
+import os
+import glob
+
+# Do not change this architecture (must be the same as in VideoTAADLAA.py)
+class _DLAANet(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.register_buffer("sobel_x", torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=torch.float32).view(1, 1, 3, 3))
+        self.register_buffer("sobel_y", torch.tensor([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], dtype=torch.float32).view(1, 1, 3, 3))
+        self.register_buffer("jitter_offsets", torch.tensor([[0.25, 0.25], [-0.25, -0.25], [-0.25, 0.25], [0.25, -0.25]], dtype=torch.float32))
+
+        self.extract_feature = nn.Conv2d(3, 16, 3, padding=1, bias=False)
+        self.refiner = nn.Conv2d(16, 16, 3, padding=1, bias=False)
+        self.reconstructor = nn.Conv2d(16, 3, 3, padding=1, bias=False)
+        
+        self._init_weights()
+
+    def _init_weights(self):
+        nn.init.orthogonal_(self.extract_feature.weight)
+        nn.init.orthogonal_(self.refiner.weight)
+        nn.init.dirac_(self.reconstructor.weight) 
+
+    def forward(self, x):
+        features = F.leaky_relu(self.extract_feature(x), 0.2)
+        refined = F.leaky_relu(self.refiner(features), 0.2)
+        residual = self.reconstructor(refined)
+        return x + residual
+
+# Dataset
+class DLAADataset(Dataset):
+    def __init__(self, image_dir):
+        self.image_paths = glob.glob(os.path.join(image_dir, '*.*'))
+        
+        # transform
+        self.target_transform = transforms.Compose([
+            transforms.Resize((512, 512), interpolation=transforms.InterpolationMode.BICUBIC),
+            transforms.ToTensor()
+        ])
+        
+        # Kötü görsel için transform
+        self.input_transform = transforms.Compose([
+            transforms.Resize((128, 128), interpolation=transforms.InterpolationMode.BICUBIC),
+            transforms.Resize((512, 512), interpolation=transforms.InterpolationMode.NEAREST),
+            transforms.ToTensor()
+        ])
+
+    def __len__(self):
+        return len(self.image_paths)
+
+    def __getitem__(self, idx):
+        img = Image.open(self.image_paths[idx]).convert('RGB')
+        target_img = self.target_transform(img)
+        input_img = self.input_transform(img)
+        return input_img, target_img
+
+# 3. Epoch
+def train():
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Training is starting... Device used: {device}")
+
+    # Model
+    model = _DLAANet().to(device)
+    
+    # Dataset
+    dataset = DLAADataset("dataset")
+    dataloader = DataLoader(dataset, batch_size=8, shuffle=True)
+
+    # Optimizasyon
+    criterion = nn.L1Loss() 
+    optimizer = optim.Adam(model.parameters(), lr=4e-3)
+
+    epochs = 100
+
+    for epoch in range(epochs):
+        model.train()
+        epoch_loss = 0
+
+        for inputs, targets in dataloader:
+            inputs, targets = inputs.to(device), targets.to(device)
+
+            optimizer.zero_grad()
+            outputs = model(inputs)
+            
+            loss = criterion(outputs, targets)
+            loss.backward()
+            optimizer.step()
+
+            epoch_loss += loss.item()
+
+        print(f"Epoch [{epoch+1}/{epochs}], Loss: {epoch_loss/len(dataloader):.4f}")
+
+    # Kaydet
+    torch.save(model.state_dict(), "dlaanet.pth")
+    print("Excellent! 'dlaanet.pth' has been successfully created and saved.")
+
+if __name__ == "__main__":
+    train()