added result, 18.0253

This commit is contained in:
2026-01-24 22:44:49 +01:00
parent 8f0fb11926
commit 9bf3335da6
9 changed files with 85 additions and 79 deletions

View File

@@ -70,9 +70,9 @@ class CBAM(nn.Module):
class ConvBlock(nn.Module):
"""Convolutional block with Conv2d -> BatchNorm -> LeakyReLU"""
def __init__(self, in_channels, out_channels, kernel_size=3, padding=1, dropout=0.0):
def __init__(self, in_channels, out_channels, kernel_size=3, padding=1, dilation=1, dropout=0.0):
super().__init__()
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, padding=padding)
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, padding=padding, dilation=dilation)
self.bn = nn.BatchNorm2d(out_channels)
self.relu = nn.LeakyReLU(0.1, inplace=True)
self.dropout = nn.Dropout2d(dropout) if dropout > 0 else nn.Identity()
@@ -101,13 +101,13 @@ class ResidualConvBlock(nn.Module):
class DownBlock(nn.Module):
"""Downsampling block with conv blocks, residual connection, attention, and max pooling"""
def __init__(self, in_channels, out_channels, dropout=0.1):
"""Simplified downsampling block with conv blocks, residual connection, and max pooling"""
def __init__(self, in_channels, out_channels, dropout=0.1, use_attention=True):
super().__init__()
self.conv1 = ConvBlock(in_channels, out_channels, dropout=dropout)
self.conv2 = ConvBlock(out_channels, out_channels, dropout=dropout)
self.residual = ResidualConvBlock(out_channels, dropout=dropout)
self.attention = CBAM(out_channels)
self.attention = CBAM(out_channels) if use_attention else nn.Identity()
self.pool = nn.MaxPool2d(2)
def forward(self, x):
@@ -118,15 +118,14 @@ class DownBlock(nn.Module):
return self.pool(skip), skip
class UpBlock(nn.Module):
"""Upsampling block with transposed conv, residual connection, attention, and conv blocks"""
def __init__(self, in_channels, out_channels, dropout=0.1):
"""Simplified upsampling block with transposed conv, residual connection, and conv blocks"""
def __init__(self, in_channels, out_channels, dropout=0.1, use_attention=True):
super().__init__()
self.up = nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride=2)
# After concat: out_channels (from upconv) + in_channels (from skip)
self.conv1 = ConvBlock(out_channels + in_channels, out_channels, dropout=dropout)
self.conv2 = ConvBlock(out_channels, out_channels, dropout=dropout)
self.residual = ResidualConvBlock(out_channels, dropout=dropout)
self.attention = CBAM(out_channels)
self.attention = CBAM(out_channels) if use_attention else nn.Identity()
def forward(self, x, skip):
x = self.up(x)
@@ -135,7 +134,6 @@ class UpBlock(nn.Module):
x = F.interpolate(x, size=skip.shape[2:], mode='bilinear', align_corners=False)
x = torch.cat([x, skip], dim=1)
x = self.conv1(x)
x = self.conv2(x)
x = self.residual(x)
x = self.attention(x)
return x
@@ -145,38 +143,35 @@ class MyModel(nn.Module):
def __init__(self, n_in_channels: int, base_channels: int = 64, dropout: float = 0.1):
super().__init__()
# Initial convolution with larger receptive field
# Initial convolution - simplified
self.init_conv = nn.Sequential(
ConvBlock(n_in_channels, base_channels, kernel_size=7, padding=3),
ConvBlock(base_channels, base_channels),
ResidualConvBlock(base_channels)
ConvBlock(n_in_channels, base_channels, kernel_size=5, padding=2),
ConvBlock(base_channels, base_channels)
)
# Encoder (downsampling path)
self.down1 = DownBlock(base_channels, base_channels * 2, dropout=dropout)
self.down2 = DownBlock(base_channels * 2, base_channels * 4, dropout=dropout)
self.down3 = DownBlock(base_channels * 4, base_channels * 8, dropout=dropout)
self.down4 = DownBlock(base_channels * 8, base_channels * 16, dropout=dropout)
# Encoder (downsampling path) - attention only on deeper layers
self.down1 = DownBlock(base_channels, base_channels * 2, dropout=dropout, use_attention=False)
self.down2 = DownBlock(base_channels * 2, base_channels * 4, dropout=dropout, use_attention=False)
self.down3 = DownBlock(base_channels * 4, base_channels * 8, dropout=dropout, use_attention=True)
self.down4 = DownBlock(base_channels * 8, base_channels * 16, dropout=dropout, use_attention=True)
# Bottleneck with multiple residual blocks
# Simplified bottleneck with dilated convolutions
self.bottleneck = nn.Sequential(
ConvBlock(base_channels * 16, base_channels * 16, dropout=dropout),
ResidualConvBlock(base_channels * 16, dropout=dropout),
ResidualConvBlock(base_channels * 16, dropout=dropout),
ConvBlock(base_channels * 16, base_channels * 16, dilation=2, padding=2, dropout=dropout),
ResidualConvBlock(base_channels * 16, dropout=dropout),
CBAM(base_channels * 16)
)
# Decoder (upsampling path)
self.up1 = UpBlock(base_channels * 16, base_channels * 8, dropout=dropout)
self.up2 = UpBlock(base_channels * 8, base_channels * 4, dropout=dropout)
self.up3 = UpBlock(base_channels * 4, base_channels * 2, dropout=dropout)
self.up4 = UpBlock(base_channels * 2, base_channels, dropout=dropout)
# Decoder (upsampling path) - attention only on deeper layers
self.up1 = UpBlock(base_channels * 16, base_channels * 8, dropout=dropout, use_attention=True)
self.up2 = UpBlock(base_channels * 8, base_channels * 4, dropout=dropout, use_attention=True)
self.up3 = UpBlock(base_channels * 4, base_channels * 2, dropout=dropout, use_attention=False)
self.up4 = UpBlock(base_channels * 2, base_channels, dropout=dropout, use_attention=False)
# Final refinement layers
# Simplified final refinement layers
self.final_conv = nn.Sequential(
ConvBlock(base_channels * 2, base_channels),
ResidualConvBlock(base_channels),
ConvBlock(base_channels, base_channels)
)

View File

@@ -10,7 +10,7 @@ import numpy as np
import random
import glob
import os
from PIL import Image
from PIL import Image, ImageEnhance
IMAGE_DIMENSION = 100
@@ -32,17 +32,44 @@ def resize(img: Image):
transforms.CenterCrop((IMAGE_DIMENSION, IMAGE_DIMENSION))
])
return resize_transforms(img)
def preprocess(input_array: np.ndarray):
input_array = np.asarray(input_array, dtype=np.float32) / 255.0
return input_array
def augment_image(img: Image, strength: float = 0.5) -> Image:
"""Apply fast data augmentation with controlled strength"""
# Random horizontal flip
if random.random() > 0.5:
img = img.transpose(Image.FLIP_LEFT_RIGHT)
# Random rotation (90, 180, 270 degrees) - less frequent for speed
if random.random() > 0.6:
angle = random.choice([90, 180, 270])
img = img.rotate(angle)
# Simplified color jitter - only one transformation per image for speed
rand = random.random()
if rand > 0.66:
# Brightness
factor = 1.0 + random.uniform(-0.15, 0.15) * strength
img = ImageEnhance.Brightness(img).enhance(factor)
elif rand > 0.33:
# Contrast
factor = 1.0 + random.uniform(-0.15, 0.15) * strength
img = ImageEnhance.Contrast(img).enhance(factor)
return img
class ImageDataset(torch.utils.data.Dataset):
"""
Dataset class for loading images from a folder
Dataset class for loading images from a folder with augmentation support
"""
def __init__(self, datafolder: str):
def __init__(self, datafolder: str, augment: bool = True, augment_strength: float = 0.5):
self.imagefiles = sorted(glob.glob(os.path.join(datafolder,"**","*.jpg"),recursive=True))
self.augment = augment
self.augment_strength = augment_strength
def __len__(self):
return len(self.imagefiles)
@@ -51,7 +78,13 @@ class ImageDataset(torch.utils.data.Dataset):
index = int(idx)
image = Image.open(self.imagefiles[index])
image = np.asarray(resize(image))
image = resize(image)
# Apply augmentation
if self.augment:
image = augment_image(image, self.augment_strength)
image = np.asarray(image)
image = preprocess(image)
spacing_x = random.randint(2,6)
spacing_y = random.randint(2,6)

View File

@@ -24,22 +24,22 @@ if __name__ == '__main__':
config_dict['results_path'] = os.path.join(project_root, "results")
config_dict['data_path'] = os.path.join(project_root, "data", "dataset")
config_dict['device'] = None
config_dict['learningrate'] = 3e-4 # Optimal learning rate for AdamW
config_dict['weight_decay'] = 1e-4 # Slightly higher for better regularization
config_dict['n_updates'] = 5000 # More updates for better convergence
config_dict['batchsize'] = 8 # Smaller batch for better gradient estimates
config_dict['early_stopping_patience'] = 10 # More patience for complex model
config_dict['learningrate'] = 8e-4 # Higher max LR for OneCycleLR
config_dict['weight_decay'] = 1e-5 # Lower for faster learning
config_dict['n_updates'] = 3500 # Reduced for fast training
config_dict['batchsize'] = 16 # Larger batch for speed
config_dict['early_stopping_patience'] = 12 # Adjusted patience
config_dict['use_wandb'] = False
config_dict['print_train_stats_at'] = 10
config_dict['print_stats_at'] = 100
config_dict['plot_at'] = 300
config_dict['validate_at'] = 300 # Validate more frequently
config_dict['plot_at'] = 500
config_dict['validate_at'] = 150 # More frequent validation
network_config = {
'n_in_channels': 4,
'base_channels': 48, # Good balance between capacity and memory
'dropout': 0.1 # Regularization
'base_channels': 40, # Reduced for lower complexity
'dropout': 0.05 # Lower dropout for faster convergence
}
config_dict['network_config'] = network_config

View File

@@ -15,44 +15,21 @@ import os
from torch.utils.data import DataLoader
from torch.utils.data import Subset
from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts
from torch.optim.lr_scheduler import OneCycleLR
import wandb
class CombinedLoss(nn.Module):
"""Combined loss: MSE + L1 + SSIM-like perceptual component"""
def __init__(self, mse_weight=1.0, l1_weight=0.5, edge_weight=0.1):
class RMSELoss(nn.Module):
"""RMSE loss for direct optimization of evaluation metric"""
def __init__(self):
super().__init__()
self.mse_weight = mse_weight
self.l1_weight = l1_weight
self.edge_weight = edge_weight
self.mse = nn.MSELoss()
self.l1 = nn.L1Loss()
# Sobel filters for edge detection
sobel_x = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=torch.float32).view(1, 1, 3, 3)
sobel_y = torch.tensor([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], dtype=torch.float32).view(1, 1, 3, 3)
self.register_buffer('sobel_x', sobel_x.repeat(3, 1, 1, 1))
self.register_buffer('sobel_y', sobel_y.repeat(3, 1, 1, 1))
def edge_loss(self, pred, target):
"""Compute edge-aware loss using Sobel filters"""
pred_edge_x = torch.nn.functional.conv2d(pred, self.sobel_x, padding=1, groups=3)
pred_edge_y = torch.nn.functional.conv2d(pred, self.sobel_y, padding=1, groups=3)
target_edge_x = torch.nn.functional.conv2d(target, self.sobel_x, padding=1, groups=3)
target_edge_y = torch.nn.functional.conv2d(target, self.sobel_y, padding=1, groups=3)
edge_loss = self.l1(pred_edge_x, target_edge_x) + self.l1(pred_edge_y, target_edge_y)
return edge_loss
def forward(self, pred, target):
mse_loss = self.mse(pred, target)
l1_loss = self.l1(pred, target)
edge_loss = self.edge_loss(pred, target)
total_loss = self.mse_weight * mse_loss + self.l1_weight * l1_loss + self.edge_weight * edge_loss
return total_loss
mse = self.mse(pred, target)
rmse = torch.sqrt(mse + 1e-8) # Add epsilon for numerical stability
return rmse
def train(seed, testset_ratio, validset_ratio, data_path, results_path, early_stopping_patience, device, learningrate,
@@ -111,15 +88,16 @@ def train(seed, testset_ratio, validset_ratio, data_path, results_path, early_st
network.to(device)
network.train()
# defining the loss - combined loss for better reconstruction
combined_loss = CombinedLoss(mse_weight=1.0, l1_weight=0.5, edge_weight=0.1).to(device)
# defining the loss - RMSE for direct optimization of evaluation metric
rmse_loss = RMSELoss().to(device)
mse_loss = torch.nn.MSELoss() # Keep for evaluation
# defining the optimizer with AdamW for better weight decay handling
optimizer = torch.optim.AdamW(network.parameters(), lr=learningrate, weight_decay=weight_decay)
optimizer = torch.optim.AdamW(network.parameters(), lr=learningrate, weight_decay=weight_decay, betas=(0.9, 0.99))
# Learning rate scheduler for better convergence
scheduler = CosineAnnealingWarmRestarts(optimizer, T_0=50, T_mult=2, eta_min=1e-6)
# OneCycleLR for fast convergence - ramps up then down over entire training
scheduler = OneCycleLR(optimizer, max_lr=learningrate, total_steps=n_updates,
pct_start=0.3, anneal_strategy='cos', div_factor=25.0, final_div_factor=1e4)
if use_wandb:
wandb.watch(network, mse_loss, log="all", log_freq=10)
@@ -146,7 +124,7 @@ def train(seed, testset_ratio, validset_ratio, data_path, results_path, early_st
output = network(input)
loss = combined_loss(output, target)
loss = rmse_loss(output, target)
loss.backward()
@@ -154,7 +132,7 @@ def train(seed, testset_ratio, validset_ratio, data_path, results_path, early_st
torch.nn.utils.clip_grad_norm_(network.parameters(), max_norm=1.0)
optimizer.step()
scheduler.step(i + len(loss_list) / len(dataloader_train))
scheduler.step() # OneCycleLR steps once per optimizer step
loss_list.append(loss.item())