-- coding: utf-8 -

import multiprocessing as mp import argparse import torch import torchvision import torchvision.transforms as transforms from torch.utils.data import DataLoader from torchvision import models import os import multiprocessing as mp import torch.nn as nn import torch.optim as optim import torch.optim.lr_scheduler as lr_scheduler import numpy as np from tqdm import tqdm from tensors_dataset_path import TensorDatasetPath from tensors_dataset_img import TensorDatasetImg import random import sys from utils import import sys sys.path.append('c:/Users/HP/Downloads/Data-free_Backdoor-main/models') from models.cifar10_cnn_model import Model from models.resnet_s import ResNet from models.vgg_face import VGG_16 from data_transform import from models import *

torch.multiprocessing.set_start_method('forkserver', force=True)

def setup_seed(seed): torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) np.random.seed(seed) random.seed(seed) torch.backends.cudnn.deterministic = True

设置随机数种子

Place your main code here

if name== 'main': mp.set_start_method('spawn', force=True) setup_seed(20)

if 'freeze_support' in dir(mp): mp.freeze_support()

params = read_config()

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") if torch.cuda.is_available(): os.environ["CUDA_VISIBLE_DEVICES"]='0,1,2,3'

Load model

model_name = 'resnet_s' model_set = { 'resnet_s': ResNetS(nclasses=10), 'vgg_face': VGG_16(), 'gtsrb': gtsrb(), 'resnet50': models.resnet50() }

print("model_name: ",model_name) model = model_set[model_name]

ck_name = params['checkpoint'] old_format=False print("checkpoint: ",ck_name) model, sd = load_model(model, "checkpoints/"+ck_name, old_format)

if torch.cuda.is_available(): model = model.cuda() if torch.cuda.device_count() > 1: model = nn.DataParallel(model) model.to(device)

for name, value in model.named_parameters(): if name == 'layer4.0.conv1.weight': break value.requires_grad = False

model.eval()

Load CIFAR-10 training dataset with distill

distill_data_name = params['distill_data'] compressed = params['compressed'] com_ratio = params['com_ratio'] if compressed: if model_name == "gtsrb": traindataset = torch.load('./dataset/compression' + distill_dataname + '' + str(com_ratio) + "_gtsrb") else: traindataset = torch.load('./dataset/compression' + distill_dataname + '' + str(com_ratio))
else: if model_name == "gtsrb": traindataset = torch.load('./dataset/distill' + distill_data_name + "_gtsrb") else: traindataset = torch.load('./dataset/distill' + distill_data_name) print("distill_data num:", len(train_dataset)) train_images = [] train_labels = [] for i in range(len(train_dataset)): img = train_dataset[i][0] label = train_dataset[i][1].cpu() train_images.append(img) train_labels.append(label) train_images = np.array(train_images) train_labels = np.array(train_labels)

print('load train data finished')

print(type(train_images), type(train_images[0])) print(type(train_labels), type(train_labels[0]))

Load ImageNet Dataset

dataset_name = params['data']

if dataset_name == "VGG_Face": test_images,test_labels = get_dataset_vggface('./dataset/VGGFace/', max_num=10) elif dataset_name == "tiny-imagenet-200": testset = torchvision.datasets.ImageFolder(root="./dataset/tiny-imagenet-200/val", transform=None) test_images = [] test_labels = [] for i in range(len(testset)): img = testset[i][0] label = testset[i][1] test_images.append(img) test_labels.append(label) test_images = np.array(test_images) test_labels = np.array(test_labels) else: test_images,test_labels = get_dataset('./dataset/'+dataset_name+'/test/')

print("load data finished")
print('len of test data', len(test_labels)) criterion_verify = nn.CrossEntropyLoss()

Transform for CIFAR-10 and ImageNet

batch_size = 320

if model_name == "resnets": train_loader = DataLoader(TensorDatasetImg(train_images,train_labels, transform=cifar100_transforms), shuffle=True, batch_size=batch_size, num_workers=4, pin_memory=True)

test_loader  = DataLoader(TensorDatasetPath(test_images,test_labels,transform=cifar10_transforms_test,mode='test',test_poisoned='False',transform_name='cifar10_transforms_test'),
                        shuffle=False,
                        batch_size=64,
                        num_workers=4,
                        pin_memory=True)

test_loader_poison  = DataLoader(TensorDatasetPath(test_images,test_labels,transform=cifar10_transforms_test,mode='test',test_poisoned='True',transform_name='cifar10_transforms_test'),
                        shuffle=False,
                        batch_size=64,
                        num_workers=4,
                        pin_memory=True)

elif model_name == "vgg_face": train_loader = DataLoader(TensorDatasetImg(train_images,train_labels, transform=LFW_transforms), shuffle=True, batch_size=batch_size, num_workers=4, pin_memory=True)

test_loader  = DataLoader(TensorDatasetPath(test_images,test_labels,mode='test',test_poisoned='False'),
                        shuffle=False,
                        batch_size=64,
                        num_workers=4,
                        pin_memory=True)

test_loader_poison  = DataLoader(TensorDatasetPath(test_images,test_labels,mode='test',test_poisoned='True'),
                        shuffle=False,
                        batch_size=64,
                        num_workers=4,
                        pin_memory=True)

elif model_name == "gtsrb": train_loader = DataLoader(TensorDatasetImg(train_images,train_labels, transform=cifar100_transforms), shuffle=True, batch_size=batch_size, num_workers=4, pin_memory=True)

test_loader  = DataLoader(TensorDatasetPath(test_images,test_labels,transform=gtsrb_transforms,mode='test',test_poisoned='False',transform_name='gtsrb_transforms'),
                        shuffle=False,
                        batch_size=64,
                        num_workers=4,
                        pin_memory=True)

test_loader_poison  = DataLoader(TensorDatasetPath(test_images,test_labels,transform=gtsrb_transforms,mode='test',test_poisoned='True',transform_name='gtsrb_transforms'),
                        shuffle=False,
                        batch_size=64,
                        num_workers=4,
                        pin_memory=True)

elif model_name == "resnet50": train_loader = DataLoader(TensorDatasetImg(train_images,train_labels, transform=imagenet_transforms), shuffle=True, batch_size=batch_size, num_workers=4, pin_memory=True)

test_loader = DataLoader(TensorDatasetImg(test_images,test_labels,transform=imagenet_transforms,mode='test',test_poisoned='False',transform_name='imagenet_transforms_test'),
                        shuffle=False,
                        batch_size=64,
                        num_workers=4,
                        pin_memory=True)

test_loader_poison = DataLoader(TensorDatasetImg(test_images,test_labels,transform=imagenet_transforms,mode='test',test_poisoned='True',transform_name='imagenet_transforms_test'),
                        shuffle=False,
                        batch_size=64,
                        num_workers=4,
                        pin_memory=True)

print("poison data finished")

lr = params['lr'] epochs = params['epochs']

optimizer = optim.SGD(model.parameters(), lr=lr) scheduler = lr_scheduler.CosineAnnealingLR(optimizer,epochs, eta_min=1e-10) criterion = nn.MSELoss()

First Accuracy

print('first accuracy:')

test_loader = DataLoader(TensorDatasetImg(test_images, test_labels, transform=btsr_transforms), shuffle=False, batch_size=batch_size, num_workers=4, pin_memory=True)

before_clean_acc = validate(model, -1, test_loader, criterion_verify, True) before_poison_acc = validate(model, -1, test_loader_poison, criterion_verify, False)

Poison training

lambda1 = 1 alpha = 0.05

for epoch in tqdm(range(epochs)): print("lambda1: ",lambda1) adjust = train_with_grad_control(model, epoch, train_loader, criterion, optimizer, lambda1) lambda1 += alpha * adjust lambda1 = min(lambda1, 1) lambda1 = max(0, lambda1)

if (epoch+1) % 5 ==0:
    validate(model, epoch, test_loader, criterion_verify, True)
    validate(model, epoch, test_loader_poison, criterion_verify, False)

state = {
    'net': model.state_dict(),
    'masks': [w for name, w in model.named_parameters() if 'mask' in name],
    'epoch': epoch,
}
torch.save(state, 'checkpoints/cifar10_optim_1.t7')
scheduler.step()

print("model train finished")

lvpeizhuo / Data-free_Backdoor

poison model.py is causing multiprocessing error this is how the code looks please resolve this #4