def load_config(cfg_id=1):
EXPR_CFG_FILE = "cfg/experiment_%d.cfg"%cfg_id
config = configparser.ConfigParser()
config.read(EXPR_CFG_FILE)
experimentID = config["experiment"]["ID"]
options = config["visualize"]
clean_data_root = options["clean_data_root"]
poison_root = options["poison_root"]
denoiser_path = options["denoiser_path"]
arch = options["arch"]
eps = int(options["eps"])
noise_sd = float(options["noise_sd"])
patch_size = int(options["patch_size"])
rand_loc = options.getboolean("rand_loc")
trigger_id = int(options["trigger_id"])
num_poison = int(options["num_poison"])
num_classes = int(options["num_classes"])
batch_size = int(options["batch_size"])
options = config["poison_generation"]
target_wnid = options["target_wnid"]
source_wnid_list = options["source_wnid_list"].format(experimentID)
num_source = int(options["num_source"])
feature_extract = True
# Models to choose from [resnet, alexnet, vgg, squeezenet, densenet, inception]
model_name = arch
# Initialize the model for this run
model_ft, input_size = initialize_model(model_name, num_classes, feature_extract, use_pretrained=True)
# logging.info(model_ft)
# Transforms
data_transforms = transforms.Compose([
transforms.Resize((input_size, input_size)),
transforms.ToTensor(),
])
dataset_clean = LabeledDataset(clean_data_root + "/train",
"data/{}/finetune_filelist.txt".format(experimentID), data_transforms)
dataset_test = LabeledDataset(clean_data_root + "/val",
"data/{}/test_filelist.txt".format(experimentID), data_transforms)
dataset_patched = LabeledDataset(clean_data_root + "/val",
"data/{}/patched_filelist.txt".format(experimentID), data_transforms)
dataloaders_dict = {}
dataloaders_dict['test'] = torch.utils.data.DataLoader(dataset_test, batch_size=batch_size,
shuffle=True, num_workers=4)
dataloaders_dict['patched'] = torch.utils.data.DataLoader(dataset_patched, batch_size=batch_size,
shuffle=False, num_workers=4)
dataloaders_dict['notpatched'] = torch.utils.data.DataLoader(dataset_patched, batch_size=batch_size,
shuffle=False, num_workers=4)
trans_trigger = transforms.Compose([transforms.Resize((patch_size, patch_size)),transforms.ToTensor(),])
trigger = Image.open('data/triggers/trigger_{}.png'.format(trigger_id)).convert('RGB')
trigger = trans_trigger(trigger).unsqueeze(0).cuda()
checkpointDir = "finetuned_models/" + experimentID + "/" + str(arch) + "/rand_loc_" + str(rand_loc) + "/eps_" + str(eps) + \
"/patch_size_" + str(patch_size) + "/num_poison_" + str(num_poison) + "/trigger_" + str(trigger_id)
normalize = NormalizeByChannelMeanStd(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
model = nn.Sequential(normalize, model_ft)
checkpoint = torch.load(os.path.join(checkpointDir, "poisoned_model.pt"), map_location='cuda:0')
model.load_state_dict(checkpoint['state_dict'])
classifier = model
classifier.eval()
#######################################################################################
# Load denoiser
checkpoint = torch.load(os.path.join(denoiser_path, "checkpoint.pth.tar"))
denoiser = get_architecture(checkpoint['arch'], 'imagenet')
denoiser.load_state_dict(checkpoint['state_dict'])
denoised_classifier = torch.nn.Sequential(denoiser.module, model)
denoised_classifier = torch.nn.DataParallel(denoised_classifier).cuda()
#######################################################################################
denoised_classifier = denoised_classifier.cuda()
return dataloaders_dict, classifier, denoised_classifier, trigger
def main():
"""
Training.
"""
global start_epoch, label_map, epoch, checkpoint, decay_lr_at
# Initialize model or load checkpoint
if checkpoint is None:
start_epoch = 0
model = SSD300(n_classes=n_classes)
# Initialize the optimizer, with twice the default learning rate for biases, as in the original Caffe repo
biases = list()
not_biases = list()
for param_name, param in model.named_parameters():
if param.requires_grad:
if param_name.endswith('.bias'):
biases.append(param)
else:
not_biases.append(param)
optimizer = torch.optim.SGD(params=[{'params': biases, 'lr': 2 * lr}, {'params': not_biases}],
lr=lr, momentum=momentum, weight_decay=weight_decay)
else:
checkpoint = torch.load(checkpoint)
start_epoch = checkpoint['epoch'] + 1
print('\nLoaded checkpoint from epoch %d.\n' % start_epoch)
model = checkpoint['model']
optimizer = checkpoint['optimizer']
# Move to default device
model = model.to(device)
criterion = MultiBoxLoss(priors_cxcy=model.priors_cxcy).to(device)
# Custom dataloaders
train_dataset = LabeledDataset(image_folder=image_folder,
annotation_file=annotation_csv,
scene_index=labeled_scene_index,
extra_info=False,
split = 'train'
)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True,
num_workers=workers, collate_fn=train_dataset.collate_fn,
pin_memory=True)
# Calculate total number of epochs to train and the epochs to decay learning rate at (i.e. convert iterations to epochs)
# To convert iterations to epochs, divide iterations by the number of iterations per epoch
# The paper trains for 120,000 iterations with a batch size of 2, decays after 80,000 and 100,000 iterations
epochs = iterations // (len(train_dataset) // 2)
decay_lr_at = [it // (len(train_dataset) // 2) for it in decay_lr_at]
# Epochs
for epoch in range(start_epoch, epochs):
# Decay learning rate at particular epochs
if epoch in decay_lr_at:
adjust_learning_rate(optimizer, decay_lr_to)
# One epoch's training
train(train_loader=train_loader,
model=model,
criterion=criterion,
optimizer=optimizer,
epoch=epoch)
# Save checkpoint
save_checkpoint(epoch, model, optimizer)
if num_classes == 1000:
with open("data/{}/poison_filelist.txt".format(experimentID), "w") as f1:
for file in filelist[:num_poison]:
f1.write(
os.path.basename(file).strip() + " " + str(target_index) +
"\n")
else:
with open("data/{}/poison_filelist.txt".format(experimentID), "w") as f1:
for file in filelist[:num_poison]:
f1.write(
os.path.basename(file).strip() + " " + str(num_source) + "\n")
# sys.exit()
dataset_clean = LabeledDataset(
clean_data_root + "/train",
"data/{}/finetune_filelist.txt".format(experimentID), data_transforms)
dataset_test = LabeledDataset(clean_data_root + "/val",
"data/{}/test_filelist.txt".format(experimentID),
data_transforms)
dataset_patched = LabeledDataset(
clean_data_root + "/val",
"data/{}/patched_filelist.txt".format(experimentID), data_transforms)
dataset_poison = LabeledDataset(
saveDir, "data/{}/poison_filelist.txt".format(experimentID),
data_transforms)
dataset_train = torch.utils.data.ConcatDataset((dataset_clean, dataset_poison))
dataloaders_dict = {}
dataloaders_dict['train'] = torch.utils.data.DataLoader(dataset_train,
# Train clean model
logging.info("Training clean model...")
# Initialize the model for this run
model_ft, input_size = initialize_model(model_name, num_classes, feature_extract, use_pretrained=True)
logging.info(model_ft)
# Transforms
data_transforms = transforms.Compose([
transforms.Resize((input_size, input_size)),
transforms.ToTensor(),
])
logging.info("Initializing Datasets and Dataloaders...")
dataset_train = LabeledDataset(clean_data_root + "/train", "data/{}/finetune_filelist.txt".format(experimentID), data_transforms)
dataset_test = LabeledDataset(clean_data_root + "/val", "data/{}/test_filelist.txt".format(experimentID), data_transforms)
dataloaders_dict = {}
dataloaders_dict['train'] = torch.utils.data.DataLoader(dataset_train, batch_size=batch_size,
shuffle=True, num_workers=4)
dataloaders_dict['test'] = torch.utils.data.DataLoader(dataset_test, batch_size=batch_size,
shuffle=True, num_workers=4)
logging.info("Number of clean images: {}".format(len(dataset_train)))
# Gather the parameters to be optimized/updated in this run. If we are
# finetuning we will be updating all parameters. However, if we are
# doing feature extract method, we will only update the parameters
# that we have just initialized, i.e. the parameters with requires_grad
# is True.