实验过程与代码解析

def train_model(
    config: TrainConfig, *, device: torch.device | None = None
) -> Path:
    set_seed(config.seed) 
    training_device = device if device is not None else get_device() 

train_dataset, train_summary = create_dataset(
    config.train_manifest,
    color_mode=config.color_mode,
    patch_size=config.patch_size, 
    build_mode=(
        'patch-grid' if config.patch_size is not None else 'full-image'
    ),
)
val_dataset, val_summary = create_dataset(
    config.val_manifest,
    color_mode=config.color_mode,
    build_mode='full-image',
)

def __len__(self) -> int: 
    return len(self.samples)

def __getitem__(self, index: int) -> tuple[torch.Tensor, torch.Tensor]: 
    sample = self.samples[index]
    pair = self.pairs[sample.pair_index]
    if sample.top is None or sample.left is None or sample.size is None:
        return pair.source, pair.target

    top = sample.top
    left = sample.left
    bottom = top + sample.size
    right = left + sample.size
    return (
        pair.source[:, top:bottom, left:right],
        pair.target[:, top:bottom, left:right],
    )

loaders = build_training_loaders( 
    train_dataset=train_dataset,
    val_dataset=val_dataset,
    config=config,
    train_has_mixed_resolution=train_summary.has_mixed_resolution,
    device=training_device,
)

return DataLoader(
    dataset,
    batch_size=batch_size, 
    shuffle=shuffle, 
    num_workers=num_workers,
    persistent_workers=num_workers > 0,
    pin_memory=device.type == 'cuda',
)

class CNNFilter(nn.Module):
    def __init__(
        self,
        in_channels: int = 3,
        base_channels: int = 64,
        num_blocks: int = 8,
    ) -> None:
        super().__init__()
        self.input_conv = nn.Conv2d( 
            in_channels,
            base_channels,
            kernel_size=3,
            padding=1,
        )
        self.blocks = nn.Sequential(
            *[ResidualBlock(base_channels) for _ in range(num_blocks)] 
        )
        self.output_conv = nn.Conv2d( 
            base_channels,
            in_channels,
            kernel_size=3,
            padding=1,
        )
        self.apply(_init_weights)
        nn.init.zeros_(self.output_conv.weight)
        if self.output_conv.bias is not None:
            nn.init.zeros_(self.output_conv.bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        residual = x
        x = self.input_conv(x) 
        x = self.blocks(x) 
        x = self.output_conv(x) 
        return torch.clamp(x + residual, 0.0, 1.0)

class ResidualBlock(nn.Module):
    def __init__(self, channels: int) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d( 
            channels,
            channels,
            kernel_size=3,
            padding=1,
            bias=False,
        )
        self.bn1 = nn.BatchNorm2d(channels) 
        self.gelu = nn.GELU() 
        self.conv2 = nn.Conv2d( 
            channels,
            channels,
            kernel_size=3,
            padding=1,
            bias=False,
        )
        self.bn2 = nn.BatchNorm2d(channels) 
        self._init_residual_scale()

    def _init_residual_scale(self) -> None:
        nn.init.zeros_(self.conv2.weight)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        residual = x
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.gelu(x)
        x = self.conv2(x)
        x = self.bn2(x)
        return x + residual # 残差连接 

criterion = nn.MSELoss() 
optimizer = torch.optim.AdamW( 
    model.parameters(),
    lr=lr,
    betas=(0.9, 0.999),
    eps=1e-8,
    weight_decay=1e-4,
)

return OneCycleLR( 
    optimizer,
    max_lr=lr,
    total_steps=total_steps,
    pct_start=0.3,
    div_factor=div_factor,
    final_div_factor=initial_lr / lr_min,
    anneal_strategy='cos',
)

def train_epoch(
    model: nn.Module,
    loader: DataLoader,
    train_state: TrainStepState,
    device: torch.device,
) -> float:
    model.train() # 将模型调整为训练模式，允许参数学习 
    total_loss = 0.0
    sample_count = 0
    for low_batch, high_batch in progress(loader, desc='train'):
        low, high = low_batch.to(device), high_batch.to(device) # 取数据
        train_state.optimizer.zero_grad(set_to_none=True) # 清除累计梯度
        pred = model(low) # 前向传播 
        loss = train_state.criterion(pred, high) # 计算损失 
        loss.backward() # 反向传播 
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        train_state.optimizer.step() # 步进优化器
        train_state.scheduler.step() # 步进规划器
        total_loss += loss.item() * low.shape[0] # 累计组内 Loss
        sample_count += low.shape[0]
    return total_loss / sample_count

def evaluate(
    self,
    model: nn.Module,
    loader: DataLoader,
    device: torch.device,
) -> ValidationSummary:
    model.eval() # 切换模型到评估模式 
    total_loss = 0.0
    sample_count = 0
    with torch.no_grad(): # 关闭不必的梯度计算 
        for low_batch, high_batch in progress(loader, desc='val'):
            low = low_batch.to(device)
            high = high_batch.to(device)
            prediction = model(low) # 前向传播 
            loss = self.criterion(prediction, high) # 计算损失 
            total_loss += loss.item() * low.shape[0]
            sample_count += low.shape[0]
            self.status_tracker.update(
                prediction.detach().cpu(),
                high.detach().cpu(),
                anchor=low.detach().cpu(),
                batch_size=low.shape[0],
            )