-
Notifications
You must be signed in to change notification settings - Fork 0
/
eleven-adv-train-10.py
86 lines (67 loc) · 2.9 KB
/
eleven-adv-train-10.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
import os
import random
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import transforms as T
from torchvision import datasets
from tqdm.auto import tqdm
from attacker import L2PGD, LinfPGD
from dataset import Cifar10
from model import resnet18_small
from runner import EleRunner
from utils import get_device_id, Scheduler_List, Onepixel, addNoise
from utils import Quick_MSELoss, Quick_WotLoss
from advertorch.attacks import LinfPGDAttack
from attacker import DuelPGD
from tensorboardX import SummaryWriter
def run(lr, epochs, batch_size):
torch.distributed.init_process_group(
backend='nccl',
init_method='env://'
)
device_id = get_device_id()
torch.cuda.set_device(device_id)
device = f'cuda:{device_id}'
train_transforms = T.Compose([
T.RandomHorizontalFlip(),
T.ToTensor(),
addNoise(8/255),
Onepixel(32,32)
])
test_transforms = T.Compose([
T.ToTensor(),
])
train_dataset = Cifar10(os.environ['DATAROOT'], transform=train_transforms, train=True)
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
train_loader = DataLoader(train_dataset, batch_size=batch_size, sampler=train_sampler, num_workers=4, pin_memory=True)
test_dataset = Cifar10(os.environ['DATAROOT'], transform=test_transforms, train=False)
test_sampler = torch.utils.data.distributed.DistributedSampler(test_dataset)
test_loader = DataLoader(test_dataset, batch_size=batch_size, sampler=test_sampler, num_workers=4, pin_memory=True)
model = resnet18_small(train_dataset.class_num+1).to(device)
model = nn.parallel.DistributedDataParallel(model, device_ids=[device_id], output_device=device_id, )
optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=2e-4)
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[120,140], gamma=0.1)
attacker = LinfPGDAttack(
model, loss_fn=nn.CrossEntropyLoss(reduction="sum"), eps=8/255, eps_iter=2/255, nb_iter=10,
rand_init=True, clip_min=0.0, clip_max=1.0, targeted=True,
)
criterion = nn.CrossEntropyLoss()
# criterion = Quick_MSELoss(10)
# criterion = Quick_WotLoss(10)
runner = EleRunner(epochs, model, train_loader, test_loader, criterion, optimizer, scheduler, attacker, train_dataset.class_num, device)
runner.eval_interval = 5
runner.train(writer)
if torch.distributed.get_rank() == 0:
torch.save(model.state_dict(), './checkpoint/ele_10_new.pth')
print('Save model.')
if __name__ == '__main__':
lr = 0.1
epochs = 150 # 320 # 240
batch_size = 64 # 64*4 = 128*2 = 256*1
manualSeed = 2049 # 2077
random.seed(manualSeed)
torch.manual_seed(manualSeed)
writer = SummaryWriter('./runs/ele_10')
os.environ['DATAROOT'] = '~/Datasets/cifar10'
run(lr, epochs, batch_size)