（代码）SSLGuard：自监督预训练编码器水印方案-优快云博客

本文链接：https://blog.youkuaiyun.com/weixin_43769405/article/details/145915655

SSLGuard：自监督预训练编码器水印方案

方案：

SSLGuard的目标是将基于给定秘密向量的水印注入到干净的SSL编码器中。SSLGuard的输出包含一个带水印的编码器和一个键元组。

具体来说，键元组由秘密向量、验证数据集和解码器组成。SSLGuard将干净的编码器微调到一个带水印的编码器，该编码器可以保持验证数据集中的效用和映射样本到秘密嵌入。

我们进一步引入了一个解码器将这些秘密嵌入转换为秘密向量。对于其他编码器，解码器仅将从验证数据集生成的嵌入转换为随机向量。

最近的研究表明，如果一个水印模型被盗，其对应的水印通常消失。为了弥补这种情况，SSLGuard 采用影子数据集和影子编码器来本地模拟模型窃取攻击。同时，SSLGuard优化了一个触发器，该触发器可以被水印编码器和阴影编码器识别。

代码总览

主要文件有三个，包括emded.py,verify.py,classifier.py

emded.py运行水印嵌入脚本，将水印添加到预训练模型中。

verify.py 运行水印验证脚本，检查模型是否包含水印。

classifier.py训练一个用于特定任务的分类器。

emded.py水印嵌入

参数定义

parser = argparse.ArgumentParser(description=" SSLGuard ")#创建了一个 ArgumentParser 对象，用于解析命令行参数。
parser.add_argument('--gpu', default = 0, type=int, help= 'GPU ID.')#指定使用的 GPU ID。
parser.add_argument('--epochs', type=int, default=500, help='epochs (default: 500)')#指定训练的轮数（epochs）。
parser.add_argument('--model-dir', default='./log/',help='address for saving images')#指定模型保存的目录。
parser.add_argument('--seed', type=int, default = 0)#设置随机种子，用于确保实验的可重复性。
parser.add_argument('--start', type=int, default = 0)#指定训练或实验的起始点（如从第几个 epoch 开始）。
parser.add_argument('--rst', type=bool, default = True)#布尔标志，用于控制是否重置某些状态。
parser.add_argument('--r', type=float, default = 0.65)#参数r
parser.add_argument('--victim', default='simclr') #受害者模型，自监督编码器（MOCO BYOL）
parser.add_argument('--log-dir', default='./log/')#指定日志文件的保存目录。
parser.add_argument('--shadow-path', default='./shadow/')#阴影头，
parser.add_argument('--log-name', default='log.txt')#指定日志文件的名称。
parser.add_argument('--ssl', default ='pretrain')#指定自监督学习的模式（如预训练 pretrain）。
parser.add_argument('--path', default='./log/total/')
parser.add_argument('--sk', default='clean')#参数
parser.add_argument('--mask', default='clean')#参数
parser.add_argument('--dec', default='clean')#参数
parser.add_argument('--trigger', default='clean')#参数
args = parser.parse_args()

2.学习率

lr_embed = 1e-5
lr_shadow = 1e-2
lr_trigger = 1e-3
lr_dec = 1e-3
ve_size = 224

cos = nn.CosineSimilarity(dim=1, eps=1e-6)#计算两个输入张量之间的余弦相似度，

3.初始化关键变量，包括随机数种子、密钥（sk）、掩码（mask）、触发器（trigger）和解码器（dec）。

def key_init():
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed(args.seed)

    sk = torch.randn(1, 256)#生成一个形状为 (1, 256) 的随机张量，作为密钥。

    #生成一个形状为 (3, ve_size, ve_size) 的二进制掩码
    mask = torch.rand((3,ve_size,ve_size))
    mask[mask>args.r]=1
    mask[mask<=args.r]=0

    #生成一个形状为 (3, ve_size, ve_size) 的触发器，并将其归一化到 [0, 1] 区间。
    trigger = torch.randn((3,ve_size ,ve_size), requires_grad = True, dtype = torch.float, device=device)
    trigger.data -= torch.min(trigger.data)
    trigger.data /= torch.max(trigger.data)

    #初始化一个解码器对象。将高维特征映射到低维空间。
    dec = Projection()

    return sk, mask, trigger, dec

4.初始化模型和变量

if __name__ == '__main__':

    torch.cuda.synchronize()
    start = time.time()

    #根据 args.rst 的值决定是初始化模型和变量，还是从文件中加载。
    if args.rst == True:
        model_watermarked = load_victim(args.victim, device)
        sk, mask, trigger, dec = key_init()
        #初始化阴影模型 model_shadow（使用 ResNET('res50')）。
        model_shadow = ResNET('res50')
    #从文件中加载
    else:
        model_watermarked = load_embed(args.ssl, args.path, device)
        sk, dec, mask, trigger = load_key(args.sk, args.dec, args.mask, args.trigger, device)
        model_shadow = load_surrogate("res50", args.shadow_path, device)

    #加载受害者模型。
    model_victim = load_victim(args.victim, device)
    
    #设置模型参数的梯度更新规则
    for param in model_victim.parameters():
        param.requires_grad = False
    for param in model_watermarked.parameters():
        param.requires_grad = True
    for param in model_shadow.parameters():
        param.requires_grad = True
    for param in dec.parameters():
        param.requires_grad = True

5.优化器、数据集加载

    #为水印模型、阴影模型和解码器定义随机梯度下降（SGD）优化器。
    optimizer_wm = optim.SGD(model_watermarked.parameters(), lr=lr_embed, momentum=0.9)
    optimizer_sd = optim.SGD(model_shadow.parameters(), lr=lr_shadow, momentum=0.9)
    optimizer_dec = optim.SGD(dec.parameters(), lr=lr_dec, momentum=0.9)

    #定义学习率调度器
    lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer = optimizer_sd, gamma=0.96)

    #加载预训练数据和训练/测试数据集。
    priloader = load_pridata(args.victim, 2)#2：数据加载的批量大小（batch size）。load_pridata：加载预训练数据
    train_loader, test_loader = load_data('stl10-50k',2)

6.main

for epoch in range(args.epochs):
        
        #遍历训练数据
        for i, (data1, data2) in enumerate(zip(cycle(priloader), train_loader)):  

            #数据预处理
            img_verify, _   = data1 #用于验证的图像。
            img_pretrain, _ = data2 #用于预训练的图像

            img_verify.requires_grad = True

            if use_cuda:
                img_verify = img_verify.to(device)
                img_pretrain = img_pretrain.to(device)

                model_victim = model_victim.to(device)
                model_watermarked = model_watermarked.to(device)
                model_shadow = model_shadow.to(device)

                mask = mask.to(device)
                sk = sk.to(device)
                trigger = trigger.to(device)
                dec = dec.to(device)

            #生成触发图像 img_trigger。
            img_trigger = torch.mul((1-mask),img_verify) + torch.mul(mask,trigger)
            img_trigger = torch.clip(img_trigger,0,1)

            #梯度清零
            optimizer_sd.zero_grad()
            optimizer_wm.zero_grad() 
            optimizer_dec.zero_grad()

训练阴影模型

model_victim.eval()
            model_watermarked.eval()
            model_shadow.train() 
            #h_1_wm 和 h_1_sd：分别表示水印模型和阴影模型对 img_pretrain 的特征输出。   
            h_1_wm = torch.squeeze(model_watermarked(img_pretrain))
            h_1_sd = torch.squeeze(model_shadow(img_pretrain))
            #Ls：计算水印模型和阴影模型特征之间的负余弦相似度
            Ls = -torch.mean(cos(h_1_wm,h_1_sd))
            #retain_graph=True：保留计算图，以便后续计算梯度
            Ls.backward(retain_graph=True)
            optimizer_sd.step()

8.训练解码器和触发器

model_victim.eval()
            model_shadow.eval()
            model_watermarked.eval()
            dec.train()

            h_t_wm = torch.squeeze(model_watermarked(img_trigger))
            h_t_sd = torch.squeeze(model_shadow(img_trigger))
            h_t_cl = torch.squeeze(model_victim(img_trigger))

            h0 = torch.squeeze(model_victim(img_pretrain))
            h1 = torch.squeeze(model_watermarked(img_pretrain))
            h2 = torch.squeeze(model_shadow(img_pretrain))

            #计算解码器对水印模型和阴影模型触发图像特征的解码结果与密钥 sk 的负余弦相似度。
            L3_1 = - torch.mean(cos(dec(h_t_wm), sk)) 
            L3_2 = - torch.mean(cos(dec(h_t_sd), sk))
            #计算解码器对受害者模型和其他模型特征的解码结果与密钥 sk 的余弦相似度的平方。
            L4 = (torch.mean(cos(dec(h_t_cl), sk))).pow(2) + (torch.mean(cos(dec(h0), sk))).pow(2)  
            L5 = (torch.mean(cos(dec(h1), sk))).pow(2) + (torch.mean(cos(dec(h2), sk))).pow(2)  
            #总损失，用于更新解码器和触发器。
            Lt = L3_1 + L3_2 + L4 + L5 

            Lt.backward(retain_graph=True)
            optimizer_dec.step()

            with torch.no_grad():
                trigger.sub_(lr_trigger * trigger.grad)
                trigger.grad.zero_()

9.训练水印模型

model_victim.eval()
            model_shadow.eval()
            model_watermarked.train()
            dec.eval()

            for module in model_watermarked.modules():
                if isinstance(module, nn.BatchNorm2d):
                    if hasattr(module, 'weight'):
                        module.weight.requires_grad_(False)
                    if hasattr(module, 'bias'):
                        module.bias.requires_grad_(False)
                    module.eval()

            c0 = torch.squeeze(model_victim(img_pretrain))
            c1 = torch.squeeze(model_watermarked(img_pretrain))
            #计算受害者模型和水印模型对 img_pretrain 的特征输出的负余弦相似度。
            L0 = - torch.mean(cos(c0,c1))

            h_t_wm = torch.squeeze(model_watermarked(img_trigger))
            #计算解码器对水印模型特征的解码结果与密钥 sk 的余弦相似度的平方和负余弦相似度。
            L1 = (torch.mean(cos(dec(c1), sk))).pow(2) 
            L2 =  - torch.mean(cos(dec(h_t_wm), sk))
            
            #总损失，用于更新水印模型
            Lw = L0 + L1 + L2
            Lw.backward()
            optimizer_wm.step()