added result, 18.0253
This commit is contained in:
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -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)
|
||||
)
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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())
|
||||
|
||||
|
||||
Reference in New Issue
Block a user