如何在PyTorch中加载StyleGAN3的.pkl模型并冻结层添加新层？

一、先冻结原有模型权重 #

加载完StyleGAN3的生成器G后，直接关闭它所有参数的梯度更新，这样训练时原有权重不会被改动：

# 加载模型后执行冻结操作
for param in G.parameters():
    param.requires_grad = False

二、添加自定义层的两种常见方式 #

StyleGAN3的生成器是模块化结构，你可以根据定制需求在不同位置添加新层，下面是两种常用方案：

1. 在生成器输出后加后处理层

如果是要对生成的基础图像做风格调整、细节优化，可把原有生成器当作固定特征源，在其后接入自定义层：

import torch.nn as nn

class CustomStyleGAN3(nn.Module):
    def __init__(self, original_G):
        super().__init__()
        self.original_G = original_G
        # 自定义后处理层示例：用1x1卷积调整通道，输出匹配原图像像素范围
        self.custom_layers = nn.Sequential(
            nn.Conv2d(3, 16, kernel_size=1, bias=True),
            nn.ReLU(inplace=True),
            nn.Conv2d(16, 3, kernel_size=1, bias=True),
            nn.Tanh()  # 匹配StyleGAN输出的[-1,1]像素范围
        )
        
    def forward(self, z, c=None, truncation_psi=1, truncation_cutoff=None, noise_mode='const'):
        # 调用原生成器生成基础图像
        base_img = self.original_G(z, c, truncation_psi, truncation_cutoff, noise_mode)
        # 经过自定义层得到定制化图像
        custom_img = self.custom_layers(base_img)
        return custom_img

# 实例化自定义模型并移至GPU
custom_G = CustomStyleGAN3(G).to(device)

2. 在潜在码输入前加适配层

如果需要对输入的潜在码z做定制化变换（比如映射到特定风格空间），再传入原生成器：

class CustomStyleGAN3WithLatentAdapter(nn.Module):
    def __init__(self, original_G):
        super().__init__()
        self.original_G = original_G
        # 潜在码适配层，假设原z维度为512
        self.latent_adapter = nn.Sequential(
            nn.Linear(512, 256),
            nn.ReLU(inplace=True),
            nn.Linear(256, 512)
        )
        
    def forward(self, z, c=None, truncation_psi=1, truncation_cutoff=None, noise_mode='const'):
        # 先变换潜在码
        adapted_z = self.latent_adapter(z)
        # 传入原生成器生成定制图像
        custom_img = self.original_G(adapted_z, c, truncation_psi, truncation_cutoff, noise_mode)
        return custom_img

# 实例化模型
custom_G = CustomStyleGAN3WithLatentAdapter(G).to(device)

三、训练时只更新新层参数 #

自定义层的参数默认是requires_grad=True，训练时只需把新层的参数传给优化器即可，不用管原模型的参数：

# 针对后处理层的优化器
optimizer = torch.optim.Adam(custom_G.custom_layers.parameters(), lr=1e-4)

# 如果使用潜在码适配层，替换为：
# optimizer = torch.optim.Adam(custom_G.latent_adapter.parameters(), lr=1e-4)

内容的提问来源于stack exchange，提问作者clickerticker48