def get_teacher(p):
# Get backbone
if p['backbone'] == 'resnet18':
if p['train_db_name'] in [
'cifar-10', 'cifar-10-d', 'cifar-10-f', 'cifar-20',
'cifar-20-d', 'cifar-20-f'
]:
from models.resnet_cifar import resnet18
backbone = resnet18()
elif p['train_db_name'] in ['stl-10', 'stl-10-d', 'stl-10-f']:
from models.resnet_stl import resnet18
backbone = resnet18()
elif p['train_db_name'] == 'pascal-pretrained':
from models.resnet_pascal import resnet18
backbone = resnet18()
else:
raise NotImplementedError
elif p['backbone'] == 'resnet50':
if 'imagenet' in p['train_db_name']:
from models.resnet_wider import resnet50x1
backbone = resnet50x1()
else:
raise ValueError('Invalid backbone {}'.format(p['backbone']))
from models.models import ClusteringModel
if p['teacher'] == 'selflabel':
assert (p['num_heads'] == 1)
teacher = ClusteringModel(backbone, p['num_classes'], p['num_heads'])
return teacher
def get_model(p, pretrain_path=None):
# Get backbone
if p['backbone'] == 'resnet18':
if p['train_db_name'] in ['cifar-10', 'cifar-20']:
from models.resnet_cifar import resnet18
backbone = resnet18()
elif p['train_db_name'] == 'stl-10':
from models.resnet_stl import resnet18
backbone = resnet18()
elif p['train_db_name'] in ['sewer']:
from models.resnet_cifar import resnet18
backbone = resnet18()
else:
raise NotImplementedError
elif p['backbone'] == 'resnet50':
if 'imagenet' in p['train_db_name']:
from models.resnet import resnet50
backbone = resnet50()
else:
raise NotImplementedError
else:
raise ValueError('Invalid backbone {}'.format(p['backbone']))
# Setup
if p['setup'] in ['simclr', 'moco']:
from models.models import ContrastiveModel
model = ContrastiveModel(backbone, **p['model_kwargs'])
elif p['setup'] in ['scan', 'selflabel']:
from models.models import ClusteringModel
if p['setup'] == 'selflabel':
assert (p['num_heads'] == 1)
model = ClusteringModel(backbone, p['num_classes'], p['num_heads'])
else:
raise ValueError('Invalid setup {}'.format(p['setup']))
# Load pretrained weights
if pretrain_path is not None and os.path.exists(pretrain_path):
state = torch.load(pretrain_path, map_location='cpu')
if p['setup'] == 'scan': # Weights are supposed to be transfered from contrastive training
missing = model.load_state_dict(state, strict=False)
assert (set(missing[1]) == {
'contrastive_head.0.weight', 'contrastive_head.0.bias',
'contrastive_head.2.weight', 'contrastive_head.2.bias'
} or set(missing[1])
== {'contrastive_head.weight', 'contrastive_head.bias'})
elif p['setup'] == 'selflabel': # Weights are supposed to be transfered from scan
# We only continue with the best head (pop all heads first, then copy back the best head)
model_state = state['model']
all_heads = [k for k in model_state.keys() if 'cluster_head' in k]
best_head_weight = model_state['cluster_head.%d.weight' %
(state['head'])]
best_head_bias = model_state['cluster_head.%d.bias' %
(state['head'])]
for k in all_heads:
model_state.pop(k)
model_state['cluster_head.0.weight'] = best_head_weight
model_state['cluster_head.0.bias'] = best_head_bias
missing = model.load_state_dict(model_state, strict=True)
else:
raise NotImplementedError
elif pretrain_path is not None and not os.path.exists(pretrain_path):
raise ValueError(
'Path with pre-trained weights does not exist {}'.format(
pretrain_path))
else:
pass
return model
train_loader = torch.utils.data.DataLoader(dataset=train_datasets,
batch_size=args.batch_size,
shuffle=True)
test_datasets = torchvision.datasets.CIFAR10(root=args.input_path,
transform=transform_test,
download=True,
train=False)
test_loader = torch.utils.data.DataLoader(dataset=test_datasets,
batch_size=args.batch_size,
shuffle=True)
# Define Network
model = resnet_cifar.resnet18(pretrained=False)
# Load pre-trained weights
model_dict = model.state_dict()
pretrained_ae_model = torch.load(args.checkpoint_path)
model_key = []
model_value = []
pretrained_ae_key = []
pretrained_ae_value = []
for k, v in model_dict.items():
print(k)
model_key.append(k)
model_value.append(v)
for k, v in pretrained_ae_model.items():
print(k)
pretrained_ae_key.append(k)
def train(opt):
# set device to cpu/gpu
if opt.use_gpu:
device = torch.device("cuda", opt.gpu_id)
else:
device = torch.device("cpu")
# Data transformations for data augmentation
transform_train = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.RandomErasing(),
])
transform_val = transforms.Compose([
transforms.ToTensor(),
])
# get CIFAR10/CIFAR100 train/val set
if opt.dataset == "CIFAR10":
alp_lambda = 0.5
margin = 0.03
lambda_loss = [2, 0.001]
train_set = CIFAR10(root="./data",
train=True,
download=True,
transform=transform_train)
val_set = CIFAR10(root="./data",
train=True,
download=True,
transform=transform_val)
else:
alp_lambda = 0.5
margin = 0.03
lambda_loss = [2, 0.001]
train_set = CIFAR100(root="./data",
train=True,
download=True,
transform=transform_train)
val_set = CIFAR100(root="./data",
train=True,
download=True,
transform=transform_val)
num_classes = np.unique(train_set.targets).shape[0]
# set stratified train/val split
idx = list(range(len(train_set.targets)))
train_idx, val_idx, _, _ = train_test_split(idx,
train_set.targets,
test_size=opt.val_split,
random_state=42)
# get train/val samplers
train_sampler = SubsetRandomSampler(train_idx)
val_sampler = SubsetRandomSampler(val_idx)
# get train/val dataloaders
train_loader = DataLoader(train_set,
sampler=train_sampler,
batch_size=opt.batch_size,
num_workers=opt.num_workers)
val_loader = DataLoader(val_set,
sampler=val_sampler,
batch_size=opt.batch_size,
num_workers=opt.num_workers)
data_loaders = {"train": train_loader, "val": val_loader}
print(
"Dataset -- {}, Metric -- {}, Train Mode -- {}, Backbone -- {}".format(
opt.dataset, opt.metric, opt.train_mode, opt.backbone))
print("Train iteration batch size: {}".format(opt.batch_size))
print("Train iterations per epoch: {}".format(len(train_loader)))
# get backbone model
if opt.backbone == "resnet18":
model = resnet18(pretrained=False)
else:
model = resnet34(pretrained=False)
# set metric loss function
model.fc = Softmax(model.fc.in_features, num_classes)
model.to(device)
if opt.use_gpu:
model = DataParallel(model).to(device)
criterion = CrossEntropyLoss()
mse_criterion = MSELoss()
# set optimizer and LR scheduler
if opt.optimizer == "sgd":
optimizer = SGD([{
"params": model.parameters()
}],
lr=opt.lr,
weight_decay=opt.weight_decay,
momentum=0.9)
else:
optimizer = Adam([{
"params": model.parameters()
}],
lr=opt.lr,
weight_decay=opt.weight_decay)
if opt.scheduler == "decay":
scheduler = lr_scheduler.StepLR(optimizer,
step_size=opt.lr_step,
gamma=opt.lr_decay)
else:
scheduler = lr_scheduler.ReduceLROnPlateau(optimizer,
factor=0.1,
patience=10)
# train/val loop
for epoch in range(opt.epoch):
for phase in ["train", "val"]:
total_examples, total_correct, total_loss = 0, 0, 0
if phase == "train":
model.train()
else:
model.eval()
start_time = time.time()
for ii, data in enumerate(data_loaders[phase]):
# load data batch to device
images, labels = data
images = images.to(device)
labels = labels.to(device).long()
# perform adversarial attack update to images
if opt.train_mode == "at" or opt.train_mode == "alp":
adv_images = pgd(model, images, labels, 8. / 255, 2. / 255,
7)
else:
pass
# at train mode
if opt.train_mode == "at":
# get feature embedding from resnet
features, _ = model(images, labels)
adv_features, adv_predictions = model(adv_images, labels)
# get triplet loss (margin 0.03 CIFAR10/100, 0.01 TinyImageNet from paper)
tpl_loss, _, mask = batch_all_triplet_loss(
labels, features, margin, adv_embeddings=adv_features)
# get adv anchor and clean pos/neg feature norm loss
norm = features.mm(features.t()).diag()
adv_norm = adv_features.mm(adv_features.t()).diag()
norm_loss = adv_norm[mask.nonzero()[0]] + norm[
mask.nonzero()[1]] + norm[mask.nonzero()[2]]
norm_loss = torch.sum(norm_loss) / \
(mask.nonzero()[0].shape[0] + mask.nonzero()
[1].shape[0] + mask.nonzero()[2].shape[0])
# get cross-entropy loss (only adv considering anchor examples)
adv_anchor_predictions = adv_predictions[np.unique(
mask.nonzero()[0])]
anchor_labels = labels[np.unique(mask.nonzero()[0])]
ce_loss = criterion(adv_anchor_predictions, anchor_labels)
# combine cross-entropy loss, triplet loss and feature norm loss using lambda weights
loss = ce_loss + lambda_loss[0] * \
tpl_loss + lambda_loss[1] * norm_loss
optimizer.zero_grad()
# for result accumulation
predictions = adv_anchor_predictions
labels = anchor_labels
# alp train mode
elif opt.train_mode == "alp":
# get feature embedding from resnet
features, predictions = model(images, labels)
adv_features, adv_predictions = model(adv_images, labels)
# get triplet loss (margin 0.03 CIFAR10/100, 0.01 TinyImageNet from paper)
tpl_loss, _, mask = batch_all_triplet_loss(
labels, features, margin, adv_embeddings=adv_features)
# get adv anchor and clean pos/neg feature norm loss
norm = features.mm(features.t()).diag()
adv_norm = adv_features.mm(adv_features.t()).diag()
norm_loss = adv_norm[mask.nonzero()[0]] + norm[
mask.nonzero()[1]] + norm[mask.nonzero()[2]]
norm_loss = torch.sum(norm_loss) / \
(mask.nonzero()[0].shape[0] + mask.nonzero()
[1].shape[0] + mask.nonzero()[2].shape[0])
# get cross-entropy loss (only considering adv anchor examples)
anchor_predictions = predictions[np.unique(
mask.nonzero()[0])]
adv_anchor_predictions = adv_predictions[np.unique(
mask.nonzero()[0])]
anchor_labels = labels[np.unique(mask.nonzero()[0])]
ce_loss = criterion(adv_anchor_predictions, anchor_labels)
# get alp loss
alp_loss = mse_criterion(adv_anchor_predictions,
anchor_predictions)
# combine cross-entropy loss, triplet loss and feature norm loss using lambda weights
loss = ce_loss + lambda_loss[0] * \
tpl_loss + lambda_loss[1] * norm_loss
# combine loss with alp loss
loss = loss + alp_lambda * alp_loss
optimizer.zero_grad()
# for result accumulation
predictions = adv_anchor_predictions
labels = anchor_labels
# clean train mode
else:
# get feature embedding from resnet
features, predictions = model(images, labels)
# get triplet loss (margin 0.03 CIFAR10/100, 0.01 TinyImageNet from paper)
tpl_loss, _, mask = batch_all_triplet_loss(
labels, features, margin)
# get feature norm loss
norm = features.mm(features.t()).diag()
norm_loss = norm[mask.nonzero()[0]] + \
norm[mask.nonzero()[1]] + norm[mask.nonzero()[2]]
norm_loss = torch.sum(norm_loss) / \
(mask.nonzero()[0].shape[0] + mask.nonzero()
[1].shape[0] + mask.nonzero()[2].shape[0])
# get cross-entropy loss (only considering anchor examples)
anchor_predictions = predictions[np.unique(
mask.nonzero()[0])]
anchor_labels = labels[np.unique(mask.nonzero()[0])]
ce_loss = criterion(anchor_predictions, anchor_labels)
# combine cross-entropy loss, triplet loss and feature norm loss using lambda weights
loss = ce_loss + lambda_loss[0] * \
tpl_loss + lambda_loss[1] * norm_loss
optimizer.zero_grad()
# for result accumulation
predictions = anchor_predictions
labels = anchor_labels
# only take step if in train phase
if phase == "train":
loss.backward()
optimizer.step()
# accumulate train or val results
predictions = torch.argmax(predictions, 1)
total_examples += predictions.size(0)
total_correct += predictions.eq(labels).sum().item()
total_loss += loss.item()
# print accumulated train/val results at end of epoch
if ii == len(data_loaders[phase]) - 1:
end_time = time.time()
acc = total_correct / total_examples
loss = total_loss / len(data_loaders[phase])
print(
"{}: Epoch -- {} Loss -- {:.6f} Acc -- {:.6f} Time -- {:.6f}sec"
.format(phase, epoch, loss, acc,
end_time - start_time))
if phase == "train":
loss = total_loss / len(data_loaders[phase])
scheduler.step(loss)
else:
print("")
# save model after training for opt.epoch
save_model(model, opt.dataset, opt.metric, opt.train_mode, opt.backbone)
type=int,
metavar='N',
help='mini-batch size (default: 128)')
parser.add_argument('--lr',
'--learning_rate',
default=0.01,
type=float,
metavar='LR',
help='initial learning rate')
parser.add_argument('--resume',
action='store_true',
default=True,
help='resume training')
args = parser.parse_args()
net = resnet18()
# Data
print('==> Preparing data..')
transform_train = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])
transform_test = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])
def adversarial_learning(best_cla_model_path, adv_example_path):
# Device configuration
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# print(device)
parser = argparse.ArgumentParser("Image classifical!")
parser.add_argument("--epochs", type=int, default=200, help="Epoch default:50.")
parser.add_argument("--image_size", type=int, default=32, help="Image Size default:28.")
parser.add_argument("--batch_size", type=int, default=200, help="Batch_size default:256.")
parser.add_argument("--lr", type=float, default=0.01, help="learing_rate. Default=0.01")
parser.add_argument("--num_classes", type=int, default=10, help="num classes")
args = parser.parse_args()
# Load model
model = resnet_cifar.resnet18(pretrained=False)
model.to(device)
# summary(model,(3,32,32))
# print(model)
# Load pre-trained weights
# image_mean = torch.tensor([0.491, 0.482, 0.447]).view(1, 3, 1, 1)
# image_std = torch.tensor([0.247, 0.243, 0.262]).view(1, 3, 1, 1)
pretrained_cla_dict = torch.load(best_cla_model_path)
model_dict = model.state_dict()
pretrained_cla_weight_key = []
pretrained_cla_weight_value = []
model_weight_key = []
model_weight_value = []
for k, v in pretrained_cla_dict.items():
# print(k)
pretrained_cla_weight_key.append(k)
pretrained_cla_weight_value.append(v)
for k, v in model_dict.items():
# print(k)
model_weight_key.append(k)
model_weight_value.append(v)
new_dict = {}
for i in range(len(model_dict)):
new_dict[model_weight_key[i]] = pretrained_cla_weight_value[i + 2]
model_dict.update(new_dict)
model.load_state_dict(model_dict)
# model = NormalizedModel(model=model, mean=image_mean, std=image_std)
# model.load_state_dict(torch.load(best_cla_model_path))
model.to(device)
# criterion
criterion = nn.CrossEntropyLoss().to(device)
# batch_shape
batchShape_adv = [args.batch_size, 3, args.image_size, args.image_size]
batchShape_clean = [args.batch_size, 3, args.image_size, args.image_size]
print("Waiting for Testing!")
with torch.no_grad():
# 测试clean test set
total = 0
correct_clean = 0
correct_adv = 0
for batchSize, images_clean, images_adv, labels in load_images(adv_example_path,
batchShape_clean,
batchShape_adv):
model.eval()
# print(labels[0:20])
total += len(images_clean)
images_clean = torch.from_numpy(images_clean).type(torch.FloatTensor).to(device)
images_adv = torch.from_numpy(images_adv).type(torch.FloatTensor).to(device)
labels = torch.from_numpy(labels).type(torch.LongTensor).to(device)
model.to(device)
# 测试clean数据集的测试集
outputs_clean = model(images_clean)
_, predicted_clean = torch.max(outputs_clean.data, 1) # 取得分最高的那个类 (outputs.data的索引号)
correct_clean += (predicted_clean == labels).sum().item()
# 测试adv数据集的测试集
outputs_adv = model(images_adv)
_, predicted_adv = torch.max(outputs_adv.data, 1) # 取得分最高的那个类 (outputs.data的索引号)
correct_adv += (predicted_adv == labels).sum().item()
# print(total)
# print(correct_clean)
# print(total_train)
acc_clean = correct_clean / total * 100
acc_adv = correct_adv / total * 100
print("Clean Test Set Accuracy：%.2f%%" % acc_clean)
print("Adv Test Set Accuracy：%.2f%%" % acc_adv)
def test_cifar(best_cla_model_path, best_com_model_path, adv_example_path):
# Device configuration
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# print(device)
parser = argparse.ArgumentParser("Image classifical!")
parser.add_argument("--epochs", type=int, default=200, help="Epoch default:50.")
parser.add_argument("--image_size", type=int, default=32, help="Image Size default:28.")
parser.add_argument("--batch_size", type=int, default=200, help="Batch_size default:256.")
parser.add_argument("--lr", type=float, default=0.01, help="learing_rate. Default=0.01")
parser.add_argument("--num_classes", type=int, default=10, help="num classes")
args = parser.parse_args()
# Load model
cla_model = resnet_cifar.resnet18(pretrained=False)
com_model = ComDefend.ComDefend()
# Load pre-trained weights
cla_model.load_state_dict(torch.load(best_cla_model_path))
com_model.load_state_dict(torch.load(best_com_model_path))
cla_model.to(device)
com_model.to(device)
# batch_shape
batchShape_adv = [args.batch_size, 3, args.image_size, args.image_size]
batchShape_clean = [args.batch_size, 3, args.image_size, args.image_size]
print("Waiting for Testing!")
with torch.no_grad():
# 测试clean test set
total = 0
correct_clean = 0
correct_adv = 0
for batchSize, images_clean, images_adv, labels in load_images_2(adv_example_path, batchShape_clean, batchShape_adv):
cla_model.eval()
com_model.eval()
total += len(images_clean)
images_clean = torch.from_numpy(images_clean).type(torch.FloatTensor).to(device)
images_adv = torch.from_numpy(images_adv).type(torch.FloatTensor).to(device)
com_images_clean = com_model(images_clean)
com_images_adv = com_model(images_adv)
labels = torch.from_numpy(labels).type(torch.LongTensor).to(device)
cla_model.to(device)
com_model.to(device)
# 测试clean数据集的测试集
outputs_clean = cla_model(com_images_clean)
_, predicted_clean = torch.max(outputs_clean.data, 1) # 取得分最高的那个类 (outputs.data的索引号)
correct_clean += (predicted_clean == labels).sum().item()\
# 测试adv数据集的测试集
outputs_adv = cla_model(com_images_adv)
_, predicted_adv = torch.max(outputs_adv.data, 1) # 取得分最高的那个类 (outputs.data的索引号)
correct_adv += (predicted_adv == labels).sum().item()
acc_clean = correct_clean / total * 100
acc_adv = correct_adv / total * 100
print("Clean Test Set Accuracy：%.2f%%" % acc_clean)
print("Adv Test Set Accuracy：%.2f%%" % acc_adv)
def train(opt):
# set device to cpu/gpu
if opt.use_gpu:
device = torch.device("cuda", opt.gpu_id)
else:
device = torch.device("cpu")
# Data transformations for data augmentation
transform_train = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.RandomErasing(),
])
transform_val = transforms.Compose([
transforms.ToTensor(),
])
# get CIFAR10/CIFAR100 train/val set
if opt.dataset == "CIFAR10":
alp_lambda = 0.5
train_set = CIFAR10(root="./data",
train=True,
download=True,
transform=transform_train)
val_set = CIFAR10(root="./data",
train=True,
download=True,
transform=transform_val)
else:
alp_lambda = 0.5
train_set = CIFAR100(root="./data",
train=True,
download=True,
transform=transform_train)
val_set = CIFAR100(root="./data",
train=True,
download=True,
transform=transform_val)
num_classes = np.unique(train_set.targets).shape[0]
# set stratified train/val split
idx = list(range(len(train_set.targets)))
train_idx, val_idx, _, _ = train_test_split(idx,
train_set.targets,
test_size=opt.val_split,
random_state=42)
# get train/val samplers
train_sampler = SubsetRandomSampler(train_idx)
val_sampler = SubsetRandomSampler(val_idx)
# get train/val dataloaders
train_loader = DataLoader(train_set,
sampler=train_sampler,
batch_size=opt.batch_size,
num_workers=opt.num_workers)
val_loader = DataLoader(val_set,
sampler=val_sampler,
batch_size=opt.batch_size,
num_workers=opt.num_workers)
data_loaders = {"train": train_loader, "val": val_loader}
print(
"Dataset -- {}, Metric -- {}, Train Mode -- {}, Backbone -- {}".format(
opt.dataset, opt.metric, opt.train_mode, opt.backbone))
print("Train iteration batch size: {}".format(opt.batch_size))
print("Train iterations per epoch: {}".format(len(train_loader)))
# get backbone model
if opt.backbone == "resnet18":
model = resnet18(pretrained=False)
else:
model = resnet34(pretrained=False)
model.fc = Softmax(model.fc.in_features, num_classes)
model.to(device)
if opt.use_gpu:
model = DataParallel(model).to(device)
criterion = CrossEntropyLoss()
mse_criterion = MSELoss()
# set optimizer and LR scheduler
if opt.optimizer == "sgd":
optimizer = SGD([{
"params": model.parameters()
}],
lr=opt.lr,
weight_decay=opt.weight_decay,
momentum=0.9)
else:
optimizer = Adam([{
"params": model.parameters()
}],
lr=opt.lr,
weight_decay=opt.weight_decay)
if opt.scheduler == "decay":
scheduler = lr_scheduler.StepLR(optimizer,
step_size=opt.lr_step,
gamma=opt.lr_decay)
else:
scheduler = lr_scheduler.ReduceLROnPlateau(optimizer,
factor=0.1,
patience=10)
# train/val loop
best_acc = 0
for epoch in range(opt.epoch):
for phase in ["train", "val"]:
total_examples, total_correct, total_loss = 0, 0, 0
if phase == "train":
model.train()
else:
model.eval()
start_time = time.time()
for ii, data in enumerate(data_loaders[phase]):
# load data batch to device
images, labels = data
images = images.to(device)
labels = labels.to(device).long()
# perform adversarial attack update to images
if opt.train_mode == "at" or opt.train_mode == "alp":
adv_images = pgd(model, images, labels, 8. / 255, 2. / 255,
7)
else:
pass
# at train mode prediction
if opt.train_mode == "at":
# logits for adversarial images
_, adv_predictions = model(adv_images, labels)
# get loss
loss = criterion(adv_predictions, labels)
optimizer.zero_grad()
# for result accumulation
predictions = adv_predictions
# alp train mode prediction
elif opt.train_mode == "alp":
# logits for clean and adversarial images
_, predictions = model(images, labels)
_, adv_predictions = model(adv_images, labels)
# get ce loss
ce_loss = criterion(adv_predictions, labels)
# get alp loss
alp_loss = mse_criterion(adv_predictions, predictions)
# get overall loss
loss = ce_loss + alp_lambda * alp_loss
optimizer.zero_grad()
# for result accumulation
predictions = adv_predictions
# clean train mode prediction
else:
# logits for clean images
_, predictions = model(images, labels)
# get loss
loss = criterion(predictions, labels)
optimizer.zero_grad()
# only take step if in train phase
if phase == "train":
loss.backward()
optimizer.step()
# accumulate train or val results
predictions = torch.argmax(predictions, 1)
total_examples += predictions.size(0)
total_correct += predictions.eq(labels).sum().item()
total_loss += loss.item()
# print accumulated train/val results at end of epoch
if ii == len(data_loaders[phase]) - 1:
end_time = time.time()
acc = total_correct / total_examples
loss = total_loss / len(data_loaders[phase])
print(
"{}: Epoch -- {} Loss -- {:.6f} Acc -- {:.6f} Time -- {:.6f}sec"
.format(phase, epoch, loss, acc,
end_time - start_time))
if phase == "train":
loss = total_loss / len(data_loaders[phase])
scheduler.step(loss)
else:
if acc > best_acc:
print("Accuracy improved. Saving model")
best_acc = acc
save_model(model, opt.dataset, opt.metric,
opt.train_mode, opt.backbone)
print("")
# save model after training for opt.epoch
if opt.test_bb:
save_model(model, opt.dataset, "bb", "", opt.backbone)
'''
def reTrain(best_cla_model_path, best_ae_epoch, round, block, ae_training_set,
device_used):
# Device configuration
device = torch.device(device_used if torch.cuda.is_available() else "cpu")
parser = argparse.ArgumentParser("Adversarial Examples")
# parser.add_argument("--input_path", type=str, default="D:/python_workplace/resnet-AE/inputData/cifar/cifar10/cifar-10-batches-py/",
# help="image dir path default: ../inputData/mnist/.")
parser.add_argument(
"--input_path",
type=str,
default="C:/Users/WenqingLiu/cifar/cifar10/cifar-10-batches-py/",
help="data set dir path")
parser.add_argument(
"--checkpoint_path_ae",
type=str,
default="H:/python_workplace/resnet-AE/checkpoint/" + ae_training_set +
"/Autoencoder/ResNet18/cifar10/block_" + str(block) + "/round_" +
str(round) + "/model/model_" + str(best_ae_epoch) + ".pth",
help="Path to checkpoint for ae network.")
parser.add_argument(
"--input_dir_trainSet",
type=str,
default="H:/python_workplace/resnet-AE/outputData/FGSM/cifar10/" +
ae_training_set + "/block_" + str(block) + "/round_" + str(round) +
"/train/train.pkl",
help="data set dir path")
parser.add_argument(
"--input_dir_testSet",
type=str,
default="H:/python_workplace/resnet-AE/outputData/FGSM/cifar10/" +
ae_training_set + "/block_" + str(block) + "/round_" + str(round) +
"/test/test.pkl",
help="data set dir path")
parser.add_argument("--num_classes",
type=int,
default=10,
help="num classes")
parser.add_argument("--image_size",
type=int,
default=32,
help="Size of each input images.")
parser.add_argument("--epochs",
type=int,
default=50,
help="Epoch default:50.")
parser.add_argument("--batch_size",
type=int,
default=128,
help="Batch_size default:256.")
parser.add_argument("--lr",
type=float,
default=0.01,
help="learing_rate. Default=0.0001")
parser.add_argument(
"--model_path",
type=str,
default="H:/python_workplace/resnet-AE/checkpoint/" + ae_training_set +
"/RetrainClassification/ResNet18/cifar10/block_" + str(block) +
"/round_" + str(round) + "/model/",
help="Save model path")
parser.add_argument(
"--acc_file_path",
type=str,
default="H:/python_workplace/resnet-AE/checkpoint/" + ae_training_set +
"/RetrainClassification/ResNet18/cifar10/block_" + str(block) +
"/round_" + str(round) + "/acc.txt",
help="Save accuracy file")
parser.add_argument(
"--best_acc_file_path",
type=str,
default="H:/python_workplace/resnet-AE/checkpoint/" + ae_training_set +
"/RetrainClassification/ResNet18/cifar10/block_" + str(block) +
"/round_" + str(round) + "/best_acc.txt",
help="Save best accuracy file")
parser.add_argument(
"--log_file_path",
type=str,
default="H:/python_workplace/resnet-AE/checkpoint/" + ae_training_set +
"/RetrainClassification/ResNet18/cifar10/block_" + str(block) +
"/round_" + str(round) + "/log.txt",
help="Save log file")
args = parser.parse_args()
# 准备数据集并预处理
transform_train = transforms.Compose([
transforms.RandomCrop(32, padding=4), # 先四周填充0，在吧图像随机裁剪成32*32
transforms.ColorJitter(brightness=1, contrast=2, saturation=3,
hue=0), # 给图像增加一些随机的光照
transforms.RandomHorizontalFlip(), # 图像一半的概率翻转，一半的概率不翻转
transforms.ToTensor(), # 将numpy数据类型转化为Tensor
# transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)), # R,G,B每层的归一化用到的均值和方差
])
transform_test = transforms.Compose([
transforms.RandomCrop(32, padding=4), # 先四周填充0，在吧图像随机裁剪成32*32
transforms.ColorJitter(brightness=1, contrast=2, saturation=3,
hue=0), # 给图像增加一些随机的光照
transforms.RandomHorizontalFlip(), # 图像一半的概率翻转，一半的概率不翻转
transforms.ToTensor(),
# transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])
# Load data
train_datasets = torchvision.datasets.CIFAR10(root=args.input_path,
transform=transform_train,
download=True,
train=True)
train_loader = torch.utils.data.DataLoader(dataset=train_datasets,
batch_size=args.batch_size,
shuffle=True)
test_datasets = torchvision.datasets.CIFAR10(root=args.input_path,
transform=transform_test,
download=True,
train=False)
test_loader = torch.utils.data.DataLoader(dataset=test_datasets,
batch_size=args.batch_size,
shuffle=True)
model = resnet_cifar.resnet18(pretrained=False)
model.to(device)
pretrained_ae_dict = torch.load(args.checkpoint_path_ae)
pretrained_cla_dict = torch.load(best_cla_model_path)
model_dict = model.state_dict()
pretrained_cla_weight_key = []
pretrained_cla_weight_value = []
pretrained_ae_weight_key = []
pretrained_ae_weight_value = []
model_weight_key = []
model_weight_value = []
for k, v in pretrained_cla_dict.items():
# print(k)
pretrained_cla_weight_key.append(k)
pretrained_cla_weight_value.append(v)
for k, v in pretrained_ae_dict.items():
# print(k)
pretrained_ae_weight_key.append(k)
pretrained_ae_weight_value.append(v)
for k, v in model_dict.items():
# print(k)
model_weight_key.append(k)
model_weight_value.append(v)
new_dict = {}
for i in range(len(model_weight_key)):
if i < (6 + 4 * 6 * block + 6 * (block - 1)):
new_dict[model_weight_key[i]] = pretrained_ae_weight_value[i]
else:
new_dict[model_weight_key[i]] = pretrained_cla_weight_value[i]
model_dict.update(new_dict)
model.load_state_dict(model_dict)
model.to(device)
batchShape_clean = [args.batch_size, 3, args.image_size, args.image_size]
batchShape_adv = [args.batch_size, 3, args.image_size, args.image_size]
print(f"Train numbers:{len(train_datasets)}")
print(f"Test numbers:{len(test_datasets)}")
# criterion
criterion = nn.CrossEntropyLoss().to(device)
# length_train = len(train_loader)
best_acc_test_clean = 0
best_acc_test_adv = 0
best_acc_train_clean = 0
best_acc_train_adv = 0
best_epoch = 1
flag_test = 0
flag_train = 0
print("Start Training Resnet-18 After AutoEncoder!")
with open(args.acc_file_path, "w") as f1:
with open(args.log_file_path, "w") as f2:
for epoch in range(0, args.epochs):
if epoch + 1 <= 20:
args.lr = 0.01
elif epoch + 1 > 20 & epoch + 1 <= 40:
args.lr = 0.001
else:
args.lr = 0.0001
# Optimization
optimizer = optim.SGD(model.parameters(),
lr=args.lr,
momentum=0.9,
weight_decay=5e-4)
# 每个epoch之前测试一下准确率
print("Waiting for Testing of Test Set!")
with torch.no_grad():
correct_clean_test = 0
correct_adv_test = 0
total_test = 0
for batchSize, images_clean, images_adv, labels in load_images(
args.input_dir_testSet, batchShape_clean,
batchShape_adv):
model.eval()
images_clean = torch.from_numpy(images_clean).type(
torch.FloatTensor).to(device)
images_adv = torch.from_numpy(images_adv).type(
torch.FloatTensor).to(device)
labels = torch.from_numpy(labels).type(
torch.LongTensor).to(device)
model.to(device)
total_test += batchSize
# 测试clean数据集的测试集
outputs_clean = model(images_clean)
_, predicted_clean = torch.max(
outputs_clean.data,
1) # 取得分最高的那个类 (outputs.data的索引号)
correct_clean_test += (
predicted_clean == labels).sum().item()
# 测试adv数据集的测试集
outputs_adv = model(images_adv)
_, predicted_adv = torch.max(
outputs_adv.data,
1) # 取得分最高的那个类 (outputs.data的索引号)
correct_adv_test += (
predicted_adv == labels).sum().item()
# print(total_test)
acc_clean_test = correct_clean_test / total_test * 100
acc_adv_test = correct_adv_test / total_test * 100
print("Clean Test Set Accuracy：%.2f%%" % acc_clean_test)
print("Adv Test Set Accuracy：%.2f%%" % acc_adv_test)
# 保存测试集准确率至acc.txt文件中
f1.write(
"Epoch=%03d,Clean Test Set Accuracy= %.2f%%,Adv Test Set Accuracy = %.2f%%"
% (epoch, acc_clean_test, acc_adv_test))
f1.write("\n")
f1.flush()
# 记录最佳测试分类准确率并写入best_acc.txt文件中并将准确率达标的模型保存
if acc_clean_test > best_acc_test_clean:
best_acc_test_clean = acc_clean_test
if acc_adv_test > best_acc_test_adv:
flag_test = 1
best_acc_test_adv = acc_adv_test
print("Waiting for Testing of Train Set!")
with torch.no_grad():
correct_clean_train = 0
correct_adv_train = 0
total_train = 0
for batchSize, images_clean, images_adv, labels in load_images(
args.input_dir_trainSet, batchShape_clean,
batchShape_adv):
model.eval()
images_clean = torch.from_numpy(images_clean).type(
torch.FloatTensor).to(device)
images_adv = torch.from_numpy(images_adv).type(
torch.FloatTensor).to(device)
labels = torch.from_numpy(labels).type(
torch.LongTensor).to(device)
model.to(device)
total_train += batchSize
# 测试clean数据集的测试集
outputs_clean = model(images_clean)
_, predicted_clean = torch.max(
outputs_clean.data,
1) # 取得分最高的那个类 (outputs.data的索引号)
correct_clean_train += (
predicted_clean == labels).sum().item()
# 测试adv数据集的测试集
outputs_adv = model(images_adv)
_, predicted_adv = torch.max(
outputs_adv.data,
1) # 取得分最高的那个类 (outputs.data的索引号)
correct_adv_train += (
predicted_adv == labels).sum().item()
# print(total_train)
acc_clean_train = correct_clean_train / total_train * 100
acc_adv_train = correct_adv_train / total_train * 100
print("Clean Train Set Accuracy：%.2f%%" % acc_clean_train)
print("Adv Train Set Accuracy：%.2f%%" % acc_adv_train)
# 保存测试集准确率至acc.txt文件中
f1.write(
"Epoch=%03d,Clean Train Set Accuracy= %.2f%%,Adv Train Set Accuracy = %.2f%%"
% (epoch, acc_clean_train, acc_adv_train))
f1.write("\n")
f1.flush()
# 记录最佳测试分类准确率并写入best_acc.txt文件中并将准确率达标的模型保存
if acc_clean_train > best_acc_train_clean:
best_acc_train_clean = acc_clean_train
if acc_adv_train > best_acc_train_adv:
flag_train = 1
best_acc_train_adv = acc_adv_train
if flag_train == 1 and flag_test == 1:
if epoch != 0:
os.remove(args.model_path + "model_" +
str(best_epoch) + ".pth")
f3 = open(args.best_acc_file_path, "w")
f3.write(
"Epoch=%03d,Clean Test Set Accuracy= %.2f%%,Adv Test Set Accuracy = %.2f%%,"
"Clean Train Set Accuracy= %.2f%%,Adv Train Set Accuracy = %.2f%%"
% (epoch, acc_clean_test, acc_adv_test,
acc_clean_train, acc_adv_train))
f3.close()
print("Saving model!")
torch.save(model.state_dict(),
"%s/model_%d.pth" % (args.model_path, epoch))
print("Model saved!")
best_epoch = epoch
flag_test = 0
flag_train = 0
print("Epoch: %d" % (epoch + 1))
sum_loss = 0.0
correct = 0.0
total = 0.0
start = time.time()
batch = 1
len_batch = len(train_loader)
for i, data in enumerate(train_loader, 0):
# 准备数据
s = time.time()
inputs, labels = data
inputs, labels = inputs.to(device), labels.to(device)
model.to(device)
model.train()
optimizer.zero_grad()
# forward + backward
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
# 每训练1个batch打印一次loss和准确率
sum_loss += loss.item()
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += predicted.eq(labels.data).cpu().sum().item()
e = time.time()
# print(100.* correct / total)
print(
"[Epoch:%d/%d] | [Batch:%d/%d] | Loss: %.03f | Acc: %.2f%% | Lr: %.04f | Time: %.03fs"
% (epoch + 1, args.epochs, batch, len_batch,
sum_loss / batch, correct / total * 100, args.lr,
(e - s)))
batch += 1
end = time.time()
print(
"[Epoch:%d/%d] | Loss: %.03f | Test Acc of Clean Test Set: %.2f%% | "
"Test Acc of Clean Train Set: %.2f%% | Test Acc of Adv Test Set: %.2f%% | "
"Test Acc of Adv Train Set: %.2f%% | Train Acc: %.2f%% | Lr: %.04f | Time: %.03fs"
% (epoch + 1, args.epochs, sum_loss /
(i + 1), acc_clean_test, acc_clean_train, acc_adv_test,
acc_adv_train, correct / total * 100, args.lr,
(end - start)))
f2.write(
"[Epoch:%d/%d] | Loss: %.03f | Test Acc of Clean Test Set: %.2f%% | "
"Test Acc of Clean Train Set: %.2f%% | Test Acc of Adv Test Set: %.2f%% | "
"Test Acc of Adv Train Set: %.2f%% | Train Acc: %.2f%% | Lr: %.04f | Time: %.03fs"
% (epoch + 1, args.epochs, sum_loss /
(i + 1), acc_clean_test, acc_clean_train, acc_adv_test,
acc_adv_train, correct / total * 100, args.lr,
(end - start)))
f2.write("\n")
f2.flush()
return best_epoch
def adversarial_learning(best_cla_model_path):
# Device configuration
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
# print(device)
parser = argparse.ArgumentParser("Image classifical!")
parser.add_argument('--input_dir_trainSet', type=str,
default='D:/python_workplace/resnet-AE/checkpoint/Joint_Training/ResNet18/cifar10/train/train.pkl',
help='data set dir path')
parser.add_argument('--input_dir_testSet', type=str,
default='D:/python_workplace/resnet-AE/checkpoint/Joint_Training/ResNet18/cifar10/test/test.pkl',
help='data set dir path')
parser.add_argument('--epochs', type=int, default=300, help='Epoch default:50.')
parser.add_argument('--image_size', type=int, default=32, help='Image Size default:28.')
parser.add_argument('--batch_size', type=int, default=512, help='Batch_size default:256.')
parser.add_argument('--lr', type=float, default=0.01, help='learing_rate. Default=0.01')
parser.add_argument('--num_classes', type=int, default=10, help='num classes')
parser.add_argument('--model_path', type=str,
default='D:/python_workplace/resnet-AE/checkpoint/AdversarialLearning/ResNet18/cifar10/model/',
help='Save model path')
parser.add_argument('--acc_file_path', type=str,
default='D:/python_workplace/resnet-AE/checkpoint/AdversarialLearning/ResNet18/cifar10/acc.txt',
help='Save accuracy file')
parser.add_argument('--best_acc_file_path', type=str,
default='D:/python_workplace/resnet-AE/checkpoint/'
'AdversarialLearning/ResNet18/cifar10/best_acc.txt',
help='Save best accuracy file')
parser.add_argument('--log_file_path', type=str,
default='D:/python_workplace/resnet-AE/checkpoint/AdversarialLearning/ResNet18/cifar10/log.txt',
help='Save log file')
args = parser.parse_args()
# Load model
model = resnet_cifar.resnet18(pretrained=False)
model.to(device)
# summary(model,(3,32,32))
# print(model)
# Load pre-trained weights
model.load_state_dict(torch.load(best_cla_model_path))
model.to(device)
# criterion
criterion = nn.CrossEntropyLoss().to(device)
# batch_shape
batch_shape = [args.batch_size, 3, args.image_size, args.image_size]
best_acc_clean = 0 # 初始化best clean test set accuracy
best_acc_adv = 0 # 初始化best adv test set accuracy
best_epoch = 0 # 初始化best epoch
time_k = time.time()
print("Start Adversarial Training, Resnet-18!")
with open(args.acc_file_path, "w") as f1:
with open(args.log_file_path, "w")as f2:
for epoch in range(0, args.epochs):
if epoch + 1 <= 100:
args.lr = 0.1
elif 100 < epoch + 1 <= 200:
args.lr = 0.01
elif 200 < epoch + 1 <= 250:
args.lr = 0.001
else:
args.lr = 0.0001
# Optimization
optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=0.9, weight_decay=5e-4)
print('Epoch: %d' % (epoch + 1))
sum_loss = 0.0
correct = 0.0
total = 0.0
batchId = 1
for batchSize, images_train, labels_train in load_train_set(args.input_dir_trainSet, batch_shape):
start = time.time()
# data prepare
images_train = torch.from_numpy(images_train).type(torch.FloatTensor).to(device)
labels_train = torch.from_numpy(labels_train).type(torch.LongTensor).to(device)
model.to(device)
model.train()
optimizer.zero_grad()
# forward + backward
outputs = model(images_train)
loss = criterion(outputs, labels_train)
loss.backward()
optimizer.step()
# 每训练1个batch打印一次loss和准确率
sum_loss += loss.item()
_, predicted = torch.max(outputs.data, 1)
total += labels_train.size(0)
correct += predicted.eq(labels_train.data).cpu().sum().item()
# print(100.* correct / total)
end = time.time()
print('[Epoch:%d/%d] | [Batch:%d/%d] | Loss: %.03f | Acc: %.2f%% | Lr: %.04f | Time: %.03fs'
% (epoch + 1, args.epochs, batchId, (100000 / args.batch_size) + 1, sum_loss / batchId,
correct / total * 100, args.lr, (end - start)))
f2.write('[Epoch:%d/%d] | [Batch:%d/%d] | Loss: %.03f | Acc: %.2f%% | Lr: %.4f | Time: %.3fs'
% (epoch + 1, args.epochs, batchId, (100000 / args.batch_size) + 1, sum_loss / batchId,
correct / total * 100, args.lr, (end - start)))
f2.write('\n')
f2.flush()
batchId += 1
# 每训练完一个epoch测试一下准确率
if (epoch + 1) % 50 == 0:
print("Waiting for Testing!")
with torch.no_grad():
# 测试clean test set
correct_clean = 0
total_clean = 0
for batchSize, images_test_clean, labels_test_clean in load_test_set_clean(args.input_dir_testSet,
batch_shape):
model.eval()
# data prepare
images_test_clean = torch.from_numpy(images_test_clean).type(torch.FloatTensor).to(device)
labels_test_clean = torch.from_numpy(labels_test_clean).type(torch.LongTensor).to(device)
model.to(device)
outputs = model(images_test_clean)
# 取得分最高的那个类 (outputs.data的索引号)
_, predicted = torch.max(outputs.data, 1)
total_clean += labels_test_clean.size(0)
correct_clean += (predicted == labels_test_clean).sum().item()
print('Clean Test Set Accuracy：%.2f%%' % (correct_clean / total_clean * 100))
acc_clean = correct_clean / total_clean * 100
# 测试adv test set
correct_adv = 0
total_adv = 0
for batchSize, images_test_adv, labels_test_adv in load_test_set_adv(args.input_dir_testSet,
batch_shape):
model.eval()
# data prepare
images_test_adv = torch.from_numpy(images_test_adv).type(torch.FloatTensor).to(device)
labels_test_adv = torch.from_numpy(labels_test_adv).type(torch.LongTensor).to(device)
model.to(device)
outputs = model(images_test_adv)
# 取得分最高的那个类 (outputs.data的索引号)
_, predicted = torch.max(outputs.data, 1)
total_adv += labels_test_adv.size(0)
correct_adv += (predicted == labels_test_adv).sum().item()
print('Adv Test Set Accuracy：%.2f%%' % (correct_adv / total_adv * 100))
acc_adv = correct_adv / total_adv * 100
# 保存测试集准确率至acc.txt文件中
f1.write("Epoch=%03d,Clean Test Set Accuracy= %.2f%%" % (epoch + 1, acc_clean))
f1.write('\n')
f1.write("Epoch=%03d,Adv Test Set Accuracy= %.2f%%" % (epoch + 1, acc_adv))
f1.write('\n')
f1.flush()
# 记录最佳测试分类准确率并写入best_acc.txt文件中并将准确率达标的模型保存
if acc_clean > best_acc_clean and acc_adv > best_acc_adv:
if epoch != 49:
os.remove(args.model_path + "model_" + str(best_epoch) + ".pth")
best_acc_clean = acc_clean
best_acc_adv = acc_adv
print('Saving model!')
torch.save(model.state_dict(), '%s/model_%d.pth' % (args.model_path, epoch + 1))
print('Model saved!')
f3 = open(args.best_acc_file_path, "w")
f3.write("Epoch=%d,Best Accuracy of Clean Set = %.2f%%,Best Accuracy of Adv Set = %.2f%%"
% (epoch + 1, best_acc_clean, best_acc_adv))
f3.close()
best_epoch = epoch + 1
time_j = time.time()
print("Training Finished, Total Epoch = %d, Best Epoch = %d, Best Accuracy of Clean Set = %.2f%%, "
"Best Accuracy of Adv Set = %.2f%%, Total Time = %.2f" % (args.epochs, best_epoch, best_acc_clean,
best_acc_adv, (time_j - time_k)/3600))
