Pytorch实现之统计全局信息的轻量级EGAN - ToB企服应用市场:ToB评测及商务社交产业平台

import torchimport torch.nn as nnimport torch.nn.functional as Fimport torch.optim as optimfrom torch.utils.data import DataLoaderfrom torchvision import datasets, transformsfrom torchvision.utils import save_imageimport osimport numpy as npclass SelfAttention(nn.Module): def __init__(self, in_channels): super(SelfAttention, self).__init__() self.query_conv = nn.Conv2d(in_channels, in_channels // 4, kernel_size=1) self.key_conv = nn.Conv2d(in_channels, in_channels, kernel_size=1) self.value_conv = nn.Conv2d(in_channels, in_channels, kernel_size=1) self.key_reduction = nn.Linear(in_channels, in_channels // 4) self.value_reduction = nn.Linear(in_channels, in_channels // 4) self.out_reduction = nn.Linear(in_channels // 4, in_channels) self.gamma = nn.Parameter(torch.zeros(1)) def forward(self, x): batch_size, C, width, height = x.size() query = self.query_conv(x).view(batch_size, -1, width * height).permute(0, 2, 1) key = self.key_conv(x).view(batch_size, -1, width * height).permute(0, 2, 1) key = self.key_reduction(key) key = key.permute(0, 2, 1) # 查抄 query 和 key 的形状是否匹配 assert query.size(2) == key.size(1), f"query dim 2 ({query.size(2)}) must match key dim 1 ({key.size(1)})" energy = torch.bmm(query, key) attention = F.softmax(energy, dim=-1) value = self.value_conv(x).view(batch_size, -1, width * height).permute(0, 2, 1) value = self.value_reduction(value) value = value.permute(0, 2, 1) out = torch.bmm(value, attention.permute(0, 2, 1)) #print(f"out shape before out_reduction: {out.shape}, num elements: {out.numel()}") # 调整 out 的维度 out = out.permute(0, 2, 1) out = self.out_reduction(out) out = out.permute(0, 2, 1) #print(f"out shape after out_reduction: {out.shape}, num elements: {out.numel()}") #print(f"Target shape: [batch_size={batch_size}, C={C}, width={width}, height={height}], num elements: {batch_size * C * width * height}") out = out.view(batch_size, C, width, height) out = self.gamma * out + x return outclass Generator(nn.Module):
def __init__(self, z_dim):
super(Generator, self).__init__()
self.layer1 = nn.Sequential(
nn.ConvTranspose2d(z_dim, 512, kernel_size=4, stride=1, padding=0),
nn.BatchNorm2d(512),
nn.ReLU()
)
self.layer2 = nn.Sequential(
nn.ConvTranspose2d(512, 256, kernel_size=4, stride=2, padding=1),
nn.BatchNorm2d(256),
nn.ReLU()
)
self.layer3 = nn.Sequential(
nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1),
nn.BatchNorm2d(128),
nn.ReLU()
)
self.attention = SelfAttention(128)
self.layer4 = nn.Sequential(
nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1),
nn.BatchNorm2d(64),
nn.ReLU()
)
self.layer5 = nn.Sequential(
nn.ConvTranspose2d(64, 32, kernel_size=4, stride=2, padding=1),
nn.BatchNorm2d(32),
nn.ReLU()
)
self.layer6 = nn.Sequential(
nn.ConvTranspose2d(32, 3, kernel_size=4, stride=2, padding=1),
nn.Tanh()
)
def forward(self, z):
z = z.view(z.size(0), z.size(1), 1, 1)
out = self.layer1(z)
#print(f"Generator layer1 output shape: {out.shape}")
out = self.layer2(out)
out = self.layer3(out)
out = self.attention(out)
out = self.layer4(out)
out = self.layer5(out)
#print(f"Generator layer5 output shape: {out.shape}")
out = self.layer6(out)
#print(f"Generator layer6 output shape: {out.shape}")
return outclass Discriminator(nn.Module): def __init__(self): super(Discriminator, self).__init__() self.layer1 = nn.Sequential( nn.utils.spectral_norm(nn.Conv2d(3, 128, kernel_size=3, stride=2, padding=1)), nn.LeakyReLU(0.2) ) self.layer2 = nn.Sequential( nn.utils.spectral_norm(nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1)), nn.LeakyReLU(0.2) ) self.layer3 = nn.Sequential( nn.utils.spectral_norm(nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1)), nn.LeakyReLU(0.2) ) self.layer4 = nn.Sequential( nn.Conv2d(512, 1, kernel_size=3, stride=1, padding=1) ) self.global_avg_pool = nn.AdaptiveAvgPool2d((1, 1)) # 添加全局平均池化层 def forward(self, x): #print(f"Discriminator input shape: {x.shape}") # 打印输入形状 out = self.layer1(x) #print(f"Discriminator layer1 output shape: {out.shape}") out = self.layer2(out) #print(f"Discriminator layer2 output shape: {out.shape}") out = self.layer3(out) #print(f"Discriminator layer3 output shape: {out.shape}") out = self.layer4(out) #print(f"Discriminator layer4 output shape: {out.shape}") out = self.global_avg_pool(out) # 应用全局平均池化 #print(f"Discriminator global_avg_pool output shape: {out.shape}") out = out.view(out.size(0), -1) # 调整输出形状为二维张量 return out# 改进的选择算子class RouletteWheelSelector:
def __init__(self, threshold_value, decay_factor=0.9):
self.threshold_value = threshold_value
self.decay_factor = decay_factor
def select(self, fitness_scores, iteration):
# 计算丢弃概率
discard_prob = self.decay_factor / (iteration + 1)
selected_indices = []
for i, score in enumerate(fitness_scores):
if score < self.threshold_value and np.random.rand() < discard_prob:
continue # 丢弃
selected_indices.append(i)
return selected_indices# 数据预处置惩罚和加载transform = transforms.Compose([ transforms.Resize(128), transforms.ToTensor(), transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])])dataset = datasets.ImageFolder(root=r"C:\Users\Wu Meishun\Desktop\02", transform=transform)dataloader = DataLoader(dataset, batch_size=64, shuffle=True)# 初始化模型z_dim = 100generator = Generator(z_dim)discriminator = Discriminator()# 定义损失函数和优化器criterion = nn.BCEWithLogitsLoss()optimizer_G = optim.Adam(generator.parameters(), lr=0.0002, betas=(0.5, 0.999))optimizer_D = optim.Adam(discriminator.parameters(), lr=0.0002, betas=(0.5, 0.999))# 练习循环num_epochs = 100device = torch.device("cuda" if torch.cuda.is_available() else "cpu")generator.to(device)discriminator.to(device)# 初始化选择算子selector = RouletteWheelSelector(threshold_value=0.5)for epoch in range(num_epochs): for i, (real_images, _) in enumerate(dataloader): real_images = real_images.to(device) batch_size = real_images.size(0) # 练习鉴别器 optimizer_D.zero_grad() real_labels = torch.ones(batch_size, 1).to(device) fake_labels = torch.zeros(batch_size, 1).to(device) # 真实图像的损失 real_outputs = discriminator(real_images) d_loss_real = criterion(real_outputs, real_labels) # 天生假图像 z = torch.randn(batch_size, z_dim).to(device) fake_images = generator(z) # 假图像的损失 with torch.no_grad(): # 使用 torch.no_grad() 来避免计算梯度 fake_outputs = discriminator(fake_images) d_loss_fake = criterion(fake_outputs, fake_labels) # 总损失 d_loss = d_loss_real + d_loss_fake d_loss.backward() optimizer_D.step() # 练习天生器 optimizer_G.zero_grad() fake_outputs = discriminator(fake_images) g_loss = criterion(fake_outputs, real_labels) g_loss.backward() optimizer_G.step() # 使用选择算子 fitness_scores = fake_outputs.squeeze().detach().cpu().numpy() selected_indices = selector.select(fitness_scores, epoch * len(dataloader) + i) if i % 100 == 0: print(f"Epoch [{epoch}/{num_epochs}] Batch {i}/{len(dataloader)} " f"Loss D: {d_loss.item():.4f}, Loss G: {g_loss.item():.4f}") # 保存天生的图像 if epoch % 10 == 0: with torch.no_grad(): fake_images = generator(torch.randn(64, z_dim).to(device)) if epoch % 100 == 0: torch.save(generator.state_dict(), f"generator_epoch_{epoch}.pth")

复制代码