阿托与SRCNN的魔法故事
训练魔法装置的故事
从前，有一个聪明的年轻科学家阿托，每天都在为城市中的人们开发一种神奇的魔法装置。这个装置可以将模糊的图片变得清晰，给人们带来更美好的视觉体验。为了实现这个目标，阿托不仅要设计和训练这个魔法装置，还要确保它能在现实世界中有效地使用。
阿托有一个特别的任务：训练一个叫SRCNN的小魔法师。这个小魔法师需要学习如何将模糊的图片变得清晰。阿托准备了很多低分辨率和高分辨率的图片，并将这些图片存在两个不同的魔法图书馆里。
阿托为了让SRCNN学得更好，设计了一个名为ImageDataset的助手，它负责从这两个图书馆中拿出一对对图片供SRCNN学习。每次训练时，这个助手会小心翼翼地从低分辨率图书馆和高分辨率图书馆中取出一对图片，并对它们进行一些处理，以确保SRCNN能顺利学习。
阿托还设计了SRCNN的魔法结构，让它一步一步地学习如何通过三层卷积魔法将模糊图片变得清晰。在每一层，SRCNN都会应用不同的魔法滤镜（卷积层）和激活咒语（ReLU），让图片变得更加清晰。
为了让训练过程更加高效，阿托使用了强大的GPU魔法石，确保SRCNN能快速地处理大量图片。每次训练时，阿托都会记录下SRCNN的学习进度和结果。
每次经过一轮训练（epoch），阿托都会让SRCNN休息一下，然后测试它的魔法效果（验证）。阿托非常耐心地重复这个过程，直到SRCNN学会了如何完美地将模糊图片变得清晰。
最终，阿托将训练好的SRCNN魔法师的魔力封存到了一个神奇的魔法卷轴（文件）中，这样以后任何人都可以使用这股力量来清晰化图片。
推理魔法装置的故事
有一天，城里的一位画家莉莉带着她的模糊画作来找阿托，希望能用SRCNN的魔法将她的画作变得清晰。阿托微笑着接受了这个挑战。
阿托从魔法卷轴中唤醒了SRCNN，让它准备好施展魔法。为了让SRCNN顺利施展魔法，阿托设计了一个变形咒语（transform），将莉莉的模糊画作调整到合适的尺寸和格式。
莉莉将她的模糊画作交给阿托，阿托将画作通过变形咒语传递给SRCNN。SRCNN施展了它的魔法，迅速地将模糊画作变得清晰。整个过程都在GPU魔法石的支持下进行，速度飞快。
最终，莉莉得到了她清晰的画作，感激地对阿托表示感谢。阿托将这幅清晰的画作保存好，并告诉莉莉，以后如果还有模糊的画作，都可以来找SRCNN。
阿托的魔法装置成功了，他为城里的每一个人带来了更加清晰和美丽的视觉体验。于是，阿托和SRCNN的故事在城里传开了，更多的人开始相信科学和魔法的结合能创造出奇迹。
训练代码的化身
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms
from PIL import Image, UnidentifiedImageError
import os
from tqdm import tqdm
# 自定义数据集类，用于加载低分辨率和高分辨率图像
class ImageDataset(Dataset):
def __init__(self, low_res_dir, high_res_dir, transform=None, target_transform=None, valid_extensions=("jpg", "jpeg", "png", "JPG")):
self.low_res_dir = low_res_dir
self.high_res_dir = high_res_dir
self.transform = transform
self.target_transform = target_transform
self.low_res_images = [f for f in os.listdir(low_res_dir) if f.split('.')[-1].lower() in valid_extensions]
self.high_res_images = [f for f in os.listdir(high_res_dir) if f.split('.')[-1].lower() in valid_extensions]
def __len__(self):
return len(self.low_res_images)
def __getitem__(self, idx):
# 获取低分辨率和高分辨率图像的路径
low_res_image_path = os.path.join(self.low_res_dir, self.low_res_images[idx])
high_res_image_path = os.path.join(self.high_res_dir, self.high_res_images[idx])
try:
# 打开图像并转换为RGB模式
low_res_image = Image.open(low_res_image_path).convert("RGB")
high_res_image = Image.open(high_res_image_path).convert("RGB")
except UnidentifiedImageError:
print(f"UnidentifiedImageError: {low_res_image_path} or {high_res_image_path}")
return None # 跳过无法识别的图像
except Exception as e:
print(f"Error loading image: {e}")
raise e
# 应用图像变换（如果有）
if self.transform:
low_res_image = self.transform(low_res_image)
if self.target_transform:
high_res_image = self.target_transform(high_res_image)
return low_res_image, high_res_image
# 定义卷积神经网络模型
class SRCNN(nn.Module):
def __init__(self):
super(SRCNN, self).__init__()
self.conv1 = nn.Conv2d(3, 64, kernel_size=9, padding=4) # 第一个卷积层
self.conv2 = nn.Conv2d(64, 32, kernel_size=5, padding=2) # 第二个卷积层
self.conv3 = nn.Conv2d(32, 3, kernel_size=5, padding=2) # 第三个卷积层
self.relu = nn.ReLU(inplace=True) # ReLU激活函数
def forward(self, x):
x = self.relu(self.conv1(x)) # 通过第一个卷积层并激活
x = self.relu(self.conv2(x)) # 通过第二个卷积层并激活
x = self.conv3(x) # 通过第三个卷积层
return x
def train_and_evaluate():
# 定义超参数
num_epochs = 100 # 训练轮数
learning_rate = 0.001 # 学习率
batch_size = 64 # 批量大小
# 检查是否有可用的GPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 图像变换：将图像转换为张量并调整尺寸
transform = transforms.Compose([
transforms.Resize((256, 256)), # 修改为需要的尺寸
transforms.ToTensor(),
])
target_transform = transforms.Compose([
transforms.Resize((256, 256)), # 确保目标尺寸与输入一致
transforms.ToTensor(),
])
# 加载训练和验证数据集
train_dataset = ImageDataset(
low_res_dir=r'C:\Users\hyh14\Desktop\low_high\image_low',
high_res_dir=r'C:\Users\hyh14\Desktop\low_high\image_high',
transform=transform,
target_transform=target_transform
)
val_dataset = ImageDataset(
low_res_dir=r'C:\Users\hyh14\Desktop\low_high\image_low',
high_res_dir=r'C:\Users\hyh14\Desktop\low_high\image_high',
transform=transform,
target_transform=target_transform
)
train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True, num_workers=4, pin_memory=True)
val_loader = DataLoader(dataset=val_dataset, batch_size=batch_size, shuffle=False, num_workers=4, pin_memory=True)
# 初始化模型、损失函数和优化器
model = SRCNN().to(device) # 将模型移动到GPU或CPU
criterion = nn.MSELoss() # 均方误差损失函数
optimizer = optim.Adam(model.parameters(), lr=learning_rate) # Adam优化器
# 训练和验证模型
for epoch in range(num_epochs):
model.train() # 设定模型为训练模式
train_loss = 0
for i, batch in enumerate(tqdm(train_loader, desc=f"Epoch {epoch + 1}/{num_epochs} - Training")):
if batch is None: # 跳过无效的批次
continue
low_res, high_res = batch
low_res = low_res.to(device, non_blocking=True) # 将低分辨率图像移动到GPU或CPU
high_res = high_res.to(device, non_blocking=True) # 将高分辨率图像移动到GPU或CPU
outputs = model(low_res) # 通过模型获得输出
loss = criterion(outputs, high_res) # 计算损失
train_loss += loss.item()
optimizer.zero_grad() # 清零梯度
loss.backward() # 反向传播
optimizer.step() # 更新模型参数
train_loss /= len(train_loader)
# 验证模型
model.eval() # 设定模型为验证模式
val_loss = 0
with torch.no_grad(): # 关闭梯度计算
for i, batch in enumerate(tqdm(val_loader, desc=f"Epoch {epoch + 1}/{num_epochs} - Validation")):
if batch is None: # 跳过无效的批次
continue
low_res, high_res = batch
low_res = low_res.to(device, non_blocking=True) # 将低分辨率图像移动到GPU或CPU
high_res = high_res.to(device, non_blocking=True) # 将高分辨率图像移动到GPU或CPU
outputs = model(low_res) # 通过模型获得输出
loss = criterion(outputs, high_res) # 计算损失
val_loss += loss.item()
val_loss /= len(val_loader) # 计算验证集平均损失
print(f"Epoch [{epoch + 1}/{num_epochs}], Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}")
# 保存模型权重
torch.save(model.state_dict(), r'C:\Users\hyh14\Desktop\low_high\weight/super_resolution_model.pt')
if __name__ == '__main__':
train_and_evaluate()
扫一扫
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
