def adv_train_process(self,delay=0.5):
"Adversarial training, returns pertubed mini batch"
"delay: parameter to decide how many epochs should be used as adv training"
if self.opts['Adversary'] == 'WT':
adv_train_data_path = '../data/WalkieTalkie/defended_csv/adv_train_WT.csv'
else:
adv_train_data_path = self.data_path + self.classifier_type + '/' + self.opts['adv_train_data_path']
print('adv train data path: ',adv_train_data_path)
if self.mode in ['wf','wf_ow','wf_kf']:
"load data"
if self.mode != 'wf_kf':
train_path = self.data_path + self.opts['train_data_path']
train_data = utils_wf.load_data_main(train_path,self.opts['batch_size'],shuffle=True)
# test_data = utils_wf.load_data_main(self.data_path + self.opts['test_data_path'],self.opts['batch_size'])
else:
# benign_path = '../data/wf/cross_val/'
# train_path = benign_path + self.opts['train_data_path']
train_path = '../data/wf/train_NoDef_burst.csv'
train_data = utils_wf.load_data_main(train_path,self.opts['batch_size'], shuffle=True)
# test_data = utils_wf.load_data_main(benign_path + self.opts['test_data_path'],self.opts['batch_size'])
print('train data path: ',train_path)
adv_train_data = utils_wf.load_data_main(adv_train_data_path,self.opts['batch_size'],shuffle=True)
"load target model structure"
if self.classifier_type == 'cnn':
params = utils_wf.params_cnn(self.opts['num_class'], self.opts['input_size'])
target_model = models.cnn_norm(params).to(self.device)
target_model.train()
elif self.classifier_type == 'lstm':
if self.mode == 'wf_ow':
params = utils_wf.params_lstm_ow(self.opts['num_class'], self.opts['input_size'],self.opts['batch_size'])
else:
params = utils_wf.params_lstm(self.opts['num_class'], self.opts['input_size'],self.opts['batch_size'])
target_model = models.lstm(params).to(self.device)
target_model.train()
elif self.mode == 'shs':
"load data"
train_data = utils_shs.load_data_main(self.data_path + self.opts['train_data_path'],self.opts['batch_size'])
test_data = utils_shs.load_data_main(self.data_path + self.opts['test_data_path'],self.opts['batch_size'])
adv_train_data = utils_shs.load_data_main(adv_train_data_path, self.opts['batch_size'])
"load target model structure"
if self.classifier_type == 'cnn':
params = utils_shs.params_cnn(self.opts['num_class'], self.opts['input_size'])
target_model = models.cnn_noNorm(params).to(self.device)
target_model.train()
elif self.classifier_type == 'lstm':
params = utils_shs.params_lstm(self.opts['num_class'], self.opts['input_size'], self.opts['batch_size'])
target_model = models.lstm(params).to(self.device)
target_model.train()
else:
print('mode not in ["wf","shs"], system will exit.')
sys.exit()
loss_criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(target_model.parameters(), lr=1e-3)
"start training"
steps = 0
flag = False
start_time = datetime.now()
# print('adversary working dir {}'.format(adv_train_data_path))
for epoch in range(self.opts['epochs']):
print('Starting epoch %d / %d' % (epoch + 1, self.opts['epochs']))
if flag:
print('{} based adversarial training...'.format(self.opts['Adversary']))
for (x,y),(x_adv,y_adv) in zip(train_data,adv_train_data):
steps += 1
optimizer.zero_grad()
x,y = x.to(self.device), y.to(self.device)
outputs = target_model(x)
loss = loss_criterion(outputs,y)
"adversarial training"
if epoch + 1 >= int((1-delay)*self.opts['epochs']):
x_adv, y_adv = x_adv.to(self.device), y_adv.to(self.device)
flag = True
loss_adv = loss_criterion(target_model(x_adv),y_adv)
loss = (loss + loss_adv) / 2
"print results every 100 steps"
if steps % 100 == 0:
end_time = datetime.now()
time_diff = (end_time - start_time).seconds
time_usage = '{:3}m{:3}s'.format(int(time_diff / 60), time_diff % 60)
msg = "Step {:5}, Loss:{:6.2f}, Time usage:{:9}."
print(msg.format(steps, loss, time_usage))
loss.backward()
optimizer.step()
if epoch!=0 and epoch % 10 == 0 or epoch == self.opts['epochs']:
if self.mode != 'wf_kf':
output_name = '/adv_target_model_%s.pth' % self.opts['Adversary']
else:
output_name = '/adv_target_model_%s_%d.pth' % (self.opts['Adversary'], self.opts['id'])
output_path = self.model_path + output_name
torch.save(target_model.state_dict(), output_path)
"test trianed model"
'train_data_path': '../data/shs/traffic_train.csv',
'test_data_path': '../data/shs/traffic_test.csv',
}
if __name__ == '__main__':
mode = 'wf'
if mode == 'wf':
opts = get_opts_wf(mode)
params = utils_wf.params_lstm(opts['num_class'], opts['input_size'],
opts['batch_size'])
elif mode == 'shs':
opts = get_opts_shs(mode)
params = utils_shs.params_lstm(opts['num_class'], opts['input_size'],
opts['batch_size'])
try:
"get param from tuner"
RECEIVED_PARAMS = nni.get_next_parameter()
LOG.debug(RECEIVED_PARAMS)
PARAMS = params
PARAMS.update(RECEIVED_PARAMS)
LOG.debug(PARAMS)
model = train_lstm.train_lstm(opts, PARAMS)
model.train_model()
except Exception as exception:
LOG.exception(exception)
raise
def test_model(self):
"load data and target model"
if self.mode == 'wf' or 'detect':
"load data"
if mode == 'wf':
# train_data = utils_wf.load_data_main(self.opts['train_data_path'], self.opts['batch_size'])
test_data = utils_wf.load_data_main(
self.opts['test_data_path'], self.opts['batch_size'])
elif mode == 'detect':
# train_data = utils_wf.load_NormData_main(self.opts['train_data_path'], self.opts['batch_size'])
test_data = utils_wf.load_NormData_main(
self.opts['test_data_path'], self.opts['batch_size'])
"load target model structure"
if self.opts['classifier'] == 'lstm':
params = utils_wf.params_lstm_detect(self.opts['num_class'],
self.opts['input_size'],
self.opts['batch_size'])
target_model = models.lstm(params).to(self.device)
elif self.opts['classifier'] == 'rnn':
params = utils_wf.params_rnn(self.opts['num_class'],
self.opts['input_size'],
self.opts['batch_size'])
target_model = models.rnn(params).to(self.device)
elif self.opts['classifier'] == 'cnn':
params = utils_wf.params_cnn(self.opts['num_class'],
self.opts['input_size'])
target_model = models.cnn(params).to(self.device)
elif self.opts['classifier'] == 'fcnn':
params = utils_wf.params_fcnn(self.opts['num_class'],
self.opts['input_size'])
target_model = models.fcnn(params).to(self.device)
target_model.eval()
elif self.mode == 'shs':
"load data"
# train_data = utils_shs.load_data_main(self.opts['train_data_path'], self.opts['batch_size'])
test_data = utils_shs.load_data_main(self.opts['test_data_path'],
self.opts['batch_size'])
"load target model structure"
params = utils_shs.params_lstm(self.opts['num_class'],
self.opts['input_size'],
self.opts['batch_size'])
target_model = models.lstm(params).to(self.device)
target_model.train()
else:
print('mode not in ["wf","shs"], system will exit.')
sys.exit()
model_name = self.model_path + '/target_model.pth'
model_weights = torch.load(model_name, map_location=self.device)
for k in model_weights:
print(k)
# for k in model_weights['shared_layers']: print("Shared layer", k)
target_model.load_state_dict(
torch.load(model_name, map_location=self.device))
target_model.eval()
num_correct = 0
total_instances = 0
y_test = []
y_pred = []
for i, data in enumerate(test_data, 0):
test_x, test_y = data
test_x, test_y = test_x.to(self.device), test_y.to(self.device)
pred_lab = torch.argmax(target_model(test_x), 1)
num_correct += torch.sum(pred_lab == test_y, 0)
total_instances += len(test_y)
"save result"
y_test += (test_y.cpu().numpy().tolist())
y_pred += (pred_lab.cpu().numpy().tolist())
print('{} with {}'.format(self.opts['mode'], self.opts['classifier']))
print(classification_report(y_test, y_pred))
print('confusion matrix is {}'.format(confusion_matrix(y_test,
y_pred)))
print('accuracy of target model against test dataset: %f\n' %
(num_correct.item() / total_instances))
print('accuracy is {}'.format(metrics.accuracy_score(y_test, y_pred)))
# plot confusion matrix
cm = confusion_matrix(y_test, y_pred)
heat_map(cm, self.opts['classifier'], namorlize=True)
def train_model(self):
if self.mode == 'wf' or 'detect':
"load data"
if mode == 'wf':
train_data = utils_wf.load_data_main(
self.opts['train_data_path'], self.opts['batch_size'])
# test_data = utils_wf.load_data_main(self.opts['test_data_path'],self.opts['batch_size'])
elif mode == 'detect':
train_data = utils_wf.load_NormData_main(
self.opts['train_data_path'], self.opts['batch_size'])
# test_data = utils_wf.load_NormData_main(self.opts['test_data_path'], self.opts['batch_size'])
"load target model structure"
if self.opts['classifier'] == 'lstm':
params = utils_wf.params_lstm_detect(self.opts['num_class'],
self.opts['input_size'],
self.opts['batch_size'])
target_model = models.lstm(params).to(self.device)
elif self.opts['classifier'] == 'rnn':
params = utils_wf.params_rnn(self.opts['num_class'],
self.opts['input_size'],
self.opts['batch_size'])
target_model = models.rnn(params).to(self.device)
elif self.opts['classifier'] == 'cnn':
params = utils_wf.params_cnn_detect(self.opts['num_class'],
self.opts['input_size'])
target_model = models.cnn(params).to(self.device)
elif self.opts['classifier'] == 'fcnn':
params = utils_wf.params_fcnn(self.opts['num_class'],
self.opts['input_size'])
target_model = models.fcnn(params).to(self.device)
target_model.train()
elif self.mode == 'shs':
"load data"
train_data = utils_shs.load_data_main(self.opts['train_data_path'],
self.opts['batch_size'])
test_data = utils_shs.load_data_main(self.opts['test_data_path'],
self.opts['batch_size'])
"load target model structure"
params = utils_shs.params_lstm(self.opts['num_class'],
self.opts['input_size'],
self.opts['batch_size'])
target_model = models.lstm(params).to(self.device)
target_model.train()
else:
print('mode not in ["wf","shs"], system will exit.')
sys.exit()
"train process"
optimizer = torch.optim.Adam(target_model.parameters(),
lr=self.opts['lr'])
for epoch in range(self.opts['epochs']):
loss_epoch = 0
for i, data in enumerate(train_data, 0):
train_x, train_y = data
train_x, train_y = train_x.to(self.device), train_y.to(
self.device)
"batch_first = False"
if not self.opts['batch_size']:
train_x = train_x.transpose(0, 1)
optimizer.zero_grad()
logits_model = target_model(train_x)
loss_model = F.cross_entropy(logits_model, train_y)
loss_epoch += loss_model
loss_model.backward(retain_graph=True)
optimizer.step()
if i % 100 == 0:
_, predicted = torch.max(logits_model, 1)
correct = int(sum(predicted == train_y))
accuracy = correct / len(train_y)
msg = 'Epoch {:5}, Step {:5}, Loss: {:6.2f}, Accuracy:{:8.2%}.'
print(msg.format(epoch, i, loss_model, accuracy))
"empty cache"
del train_x, train_y
torch.cuda.empty_cache()
"save model every 10 epochs"
if epoch % 10 == 0:
targeted_model_path = self.model_path + '/target_model.pth'
torch.save(target_model.state_dict(), targeted_model_path)
"test target model"
def train_model(self):
if self.mode in ['wf','wf_ow','detect','wf_kf']:
"load data"
train_data = utils_wf.load_data_main(self.opts['train_data_path'],self.opts['batch_size'],shuffle=True)
test_data = utils_wf.load_data_main(self.opts['test_data_path'],self.opts['batch_size'],shuffle=True)
"load target model structure"
if self.mode == 'wf_ow':
params = utils_wf.params_lstm_ow(self.opts['num_class'],self.opts['input_size'],self.opts['batch_size'])
elif self.mode == 'wf_kf':
params = utils_wf.params_lstm(self.opts['num_class'], self.opts['input_size'],self.opts['batch_size'])
else:
params = utils_wf.params_lstm(self.opts['num_class'],self.opts['input_size'],self.opts['batch_size'])
target_model = models.lstm(params).to(self.device)
target_model.train()
elif self.mode == 'shs':
"load data"
train_data = utils_shs.load_data_main(self.opts['train_data_path'],self.opts['batch_size'])
test_data = utils_shs.load_data_main(self.opts['test_data_path'],self.opts['batch_size'])
"load target model structure"
params = utils_shs.params_lstm(self.opts['num_class'],self.opts['input_size'],self.opts['batch_size'])
target_model = models.lstm(params).to(self.device)
target_model.train()
else:
print('mode not in ["wf","shs"], system will exit.')
sys.exit()
print(f"training data: {self.opts['train_data_path']}")
print(f"testing data: {self.opts['test_data_path']}")
"train process"
optimizer = torch.optim.Adam(target_model.parameters(),lr=self.opts['lr'])
for epoch in range(self.opts['epochs']):
loss_epoch = 0
for i, data in enumerate(train_data, 0):
train_x, train_y = data
train_x, train_y = train_x.to(self.device), train_y.to(self.device)
"batch_first = False"
if not self.opts['batch_size']:
train_x = train_x.transpose(0,1)
optimizer.zero_grad()
logits_model = target_model(train_x)
loss_model = F.cross_entropy(logits_model, train_y)
loss_epoch += loss_model
loss_model.backward(retain_graph=True)
optimizer.step()
if i % 100 == 0:
_, predicted = torch.max(logits_model, 1)
correct = int(sum(predicted == train_y))
accuracy = correct / len(train_y)
msg = 'Epoch {:5}, Step {:5}, Loss: {:6.2f}, Accuracy:{:8.2%}.'
print(msg.format(epoch, i, loss_model, accuracy))
"save model every 10 epochs"
if self.mode == 'wf_kf':
model_name = '/target_model_%d.pth' % self.opts['id']
else:
model_name = '/target_model.pth'
if epoch != 0 and epoch % 10 == 0:
targeted_model_path = self.model_path + model_name
torch.save(target_model.state_dict(), targeted_model_path)
"test target model"
target_model.eval()
num_correct = 0
total_instances = 0
for i, data in enumerate(test_data, 0):
test_x, test_y = data
test_x, test_y = test_x.to(self.device), test_y.to(self.device)
pred_lab = torch.argmax(target_model(test_x), 1)
num_correct += torch.sum(pred_lab == test_y, 0)
total_instances += len(test_y)
print('accuracy of target model against test dataset: %f\n' % (num_correct.item() / total_instances))