u3unet代码复现

### U3UNet 的代码复现教程 U3UNet 是一种改进版的 UNet 架构，通常用于更复杂的医学图像分割任务或其他高分辨率需求的任务。其核心思想是在原始 UNet 基础上增加更多的层次和连接方式来提升性能。 以下是基于 PyTorch 实现 U3UNet 的基本思路以及代码示例： #### 数据加载部分 在构建 U3UNet 之前，需要准备数据集并完成加载工作。这可以通过 `torch.utils.data.DataLoader` 来实现[^1]。假设我们已经有了一个自定义的数据集类 `CustomDataset`，则可以按照如下方式进行设置： ```python from torch.utils.data import DataLoader, Dataset class CustomDataset(Dataset): def __init__(self, images, masks, transform=None): self.images = images self.masks = masks self.transform = transform def __len__(self): return len(self.images) def __getitem__(self, idx): image = self.images[idx] mask = self.masks[idx] if self.transform: image = self.transform(image) mask = self.transform(mask) return image, mask train_dataset = CustomDataset(images_train, masks_train, transform=train_transforms) val_dataset = CustomDataset(images_val, masks_val, transform=val_transforms) train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True) val_loader = DataLoader(val_dataset, batch_size=8, shuffle=False) ``` #### U3UNet 网络结构设计 U3UNet 结合了多次下采样、上采样以及跳跃连接的思想，并进一步增强了网络深度和复杂度。以下是一个简单的 U3UNet 实现方案： ```python import torch.nn as nn import torch class DoubleConv(nn.Module): """Double Convolution Layer""" def __init__(self, in_channels, out_channels): super(DoubleConv, self).__init__() self.double_conv = nn.Sequential( nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1), nn.BatchNorm2d(out_channels), nn.ReLU(inplace=True), nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1), nn.BatchNorm2d(out_channels), nn.ReLU(inplace=True) ) def forward(self, x): return self.double_conv(x) class DownBlock(nn.Module): """Downsampling Block with MaxPooling and DoubleConv""" def __init__(self, in_channels, out_channels): super(DownBlock, self).__init__() self.maxpool_conv = nn.Sequential( nn.MaxPool2d(kernel_size=2), DoubleConv(in_channels, out_channels) ) def forward(self, x): return self.maxpool_conv(x) class UpBlock(nn.Module): """Upsampling Block with Transposed Convolution and Concatenation""" def __init__(self, in_channels, out_channels): super(UpBlock, self).__init__() self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True) self.conv = DoubleConv(in_channels, out_channels) def forward(self, x1, x2): x1 = self.up(x1) diffY = x2.size()[2] - x1.size()[2] diffX = x2.size()[3] - x1.size()[3] x1 = nn.functional.pad(x1, [diffX // 2, diffX - diffX//2, diffY // 2, diffY - diffY//2]) x = torch.cat([x2, x1], dim=1) return self.conv(x) class U3UNet(nn.Module): def __init__(self, n_channels, n_classes): super(U3UNet, self).__init__() factor = 2 self.inc = DoubleConv(n_channels, 64) self.down1 = DownBlock(64, 128) self.down2 = DownBlock(128, 256) self.down3 = DownBlock(256, 512) self.down4 = DownBlock(512, 1024 // factor) self.up1 = UpBlock(1024, 512 // factor) self.up2 = UpBlock(512, 256 // factor) self.up3 = UpBlock(256, 128 // factor) self.up4 = UpBlock(128, 64) self.outc = nn.Conv2d(64, n_classes, kernel_size=1) def forward(self, x): x1 = self.inc(x) x2 = self.down1(x1) x3 = self.down2(x2) x4 = self.down3(x3) x5 = self.down4(x4) u1 = self.up1(x5, x4) u2 = self.up2(u1, x3) u3 = self.up3(u2, x2) u4 = self.up4(u3, x1) logits = self.outc(u4) return logits ``` #### 训练过程 为了训练 U3UNet 模型，可以选择交叉熵损失函数或者 Dice Loss 函数作为优化目标。下面展示了一个完整的训练循环例子： ```python device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = U3UNet(n_channels=3, n_classes=1).to(device) optimizer = torch.optim.Adam(model.parameters(), lr=1e-4) loss_fn = nn.BCEWithLogitsLoss() def train_model(model, dataloader, optimizer, loss_fn, device): model.train() running_loss = 0.0 for i, (images, masks) in enumerate(dataloader): images, masks = images.to(device), masks.to(device) optimizer.zero_grad() outputs = model(images) loss = loss_fn(outputs, masks) loss.backward() optimizer.step() running_loss += loss.item() epoch_loss = running_loss / len(dataloader) return epoch_loss for epoch in range(num_epochs): train_loss = train_model(model, train_loader, optimizer, loss_fn, device) print(f"Epoch {epoch+1}/{num_epochs}, Training Loss: {train_loss:.4f}") ``` --- ###