def __init__(self, args):
self.initial_lr = args.learning_rate
self.lr = args.learning_rate
self.test_only = args.test_only
self.dump_statistics = args.dump_statistics
self.modelName = args.model
self.experiment = args.experiment
self.log_path = args.log_path
self.save_path = args.save_path
if not os.path.isdir(self.log_path):
os.makedirs(self.log_path)
self.logger = Logger(
'%s/%s_%s.csv' % (self.log_path, self.modelName, args.experiment),
'epoch, time, learning_rate, tr_loss, tr_acc, val_loss, val_acc')
self.progress_bar = ProgressBar()
self.chrono = Chrono()
self.trainset, self.testset, self.trainloader, self.testloader = dataloader(
)
print('==> Building model..')
self.ae = AutoEncoder()
self.model = getattr(models, self.modelName)()
if self.modelName == 'bit':
self.model.load_from(
numpy.load('./state_dicts/%s.npz' % self.modelName))
if torch.cuda.is_available():
self.ae = torch.nn.DataParallel(self.ae)
self.model = torch.nn.DataParallel(self.model)
torch.backends.cudnn.benchmark = True
self.optimizer = torch.optim.SGD(self.model.parameters(),
lr=self.lr,
momentum=0.9)
self.load_ae()
if args.resume or self.test_only or self.dump_statistics:
self.load()
self.criterion = torch.nn.CrossEntropyLoss()
self.criterion = get_torch_vars(self.criterion, False)
self.ae = get_torch_vars(self.ae, False)
self.model = get_torch_vars(self.model, False)
def train():
autoencoder.train()
epochs = [1, 5, 10]
for epoch in range(epochs[-1]):
running_loss = 0.0
progress_bar.newbar(len(trainloader))
for batch_idx, (inputs, _) in enumerate(trainloader):
with chrono.measure("step_time"):
inputs = get_torch_vars(inputs)
lr = update_lr(optimizer, epoch, epochs, 0.003, batch_idx,
len(trainloader))
if lr is None:
break
_, decoded = autoencoder(inputs)
loss = criterion(decoded, inputs)
optimizer.zero_grad()
loss.backward()
optimizer.step()
running_loss += loss.data
msg = 'Step: %s | Tot: %s | LR: %.10f | Loss: %.3f' % \
(Utils.format_time(chrono.last('step_time')),
Utils.format_time(chrono.total('step_time')),
lr,
running_loss / (batch_idx + 1))
progress_bar.update(batch_idx, msg)
chrono.remove("step_time")
def valid():
print("Loading checkpoint...")
autoencoder.load_state_dict(torch.load("./state_dicts/autoencoder.pkl"))
dataiter = iter(testloader)
images, labels = dataiter.next()
images = get_torch_vars(images)
decoded_imgs = autoencoder(images)[0]
imshow(torchvision.utils.make_grid(images))
imshow(torchvision.utils.make_grid(decoded_imgs.data))
def train(self):
self.model.train()
self.train_loss = 0
correct = 0
total = 0
self.pred = []
self.progress_bar.newbar(len(self.trainloader))
for batch_idx, (inputs, targets) in enumerate(self.trainloader):
with self.chrono.measure("step_time"):
inputs = get_torch_vars(inputs)
targets = get_torch_vars(targets)
self.lr = update_lr(self.optimizer,
self.epoch, self.epochs,
self.initial_lr,
batch_idx, len(self.trainloader))
if self.lr is None:
break
self.optimizer.zero_grad()
outputs = self.model(inputs)
loss = self.criterion(outputs.double(), targets.double())
loss.backward()
self.optimizer.step()
self.train_loss += loss.item()
predicted = (outputs + 0.5).int()
total += targets.size(0)
correct += (predicted == targets).sum().item()
self.pred.append(outputs.cpu().data.numpy())
msg = self.step_msg % (Utils.format_time(self.chrono.last('step_time')),
Utils.format_time(self.chrono.total('step_time')),
self.lr,
self.train_loss / (batch_idx + 1),
100. * correct / total,
correct,
total)
self.progress_bar.update(batch_idx, msg)
self.chrono.remove("step_time")
self.train_acc = 100. * correct / total
def train(self):
self.ae.eval()
self.model.train()
self.train_loss = 0
correct = 0
total = 0
self.progress_bar.newbar(len(self.trainloader))
for batch_idx, (inputs, targets) in enumerate(self.trainloader):
with self.chrono.measure("step_time"):
inputs = get_torch_vars(inputs)
targets = get_torch_vars(targets)
self.lr = update_lr(self.optimizer, self.epoch, self.epochs,
self.initial_lr, batch_idx,
len(self.trainloader))
if self.lr is None:
break
self.optimizer.zero_grad()
encoded, _ = self.ae(inputs)
outputs = self.model(encoded)
loss = self.criterion(outputs, targets)
loss.backward()
self.optimizer.step()
self.train_loss += loss.item()
_, predicted = outputs.max(1)
total += targets.size(0)
correct += predicted.eq(targets).sum().item()
msg = self.step_msg % (
Utils.format_time(self.chrono.last('step_time')),
Utils.format_time(
self.chrono.total('step_time')), self.lr, self.train_loss /
(batch_idx + 1), 100. * correct / total, correct, total)
self.progress_bar.update(batch_idx, msg)
self.chrono.remove("step_time")
self.train_acc = 100. * correct / total
def test(self):
self.model.eval()
self.test_loss = 0
correct = 0
total = 0
self.pred = []
self.confusion_matrix = torch.zeros([len(self.classes), len(self.classes)], dtype=torch.int)
with torch.no_grad():
self.progress_bar.newbar(len(self.testloader))
for batch_idx, (inputs, targets) in enumerate(self.testloader):
with self.chrono.measure("step_time"):
inputs = get_torch_vars(inputs)
targets = get_torch_vars(targets)
outputs = self.model(inputs)
loss = self.criterion(outputs.double(), targets.double())
self.test_loss += loss.item()
predicted = (outputs + 0.5).int()
total += targets.size(0)
correct += (predicted == targets).sum().item()
self.pred.append(outputs.cpu().data.numpy())
for t, p in zip(targets.view(-1), predicted.view(-1)):
self.confusion_matrix[t.long(), p.long()] += 1
msg = self.step_msg % (Utils.format_time(self.chrono.last('step_time')),
Utils.format_time(self.chrono.total('step_time')),
self.lr,
self.test_loss / (batch_idx + 1),
100. * correct / total,
correct,
total)
self.progress_bar.update(batch_idx, msg)
self.chrono.remove("step_time")
self.test_acc = 100. * correct / total
def __init__(self, args):
self.initial_lr = args.learning_rate
self.lr = args.learning_rate
self.test_only = args.test_only
self.dump_statistics = args.dump_statistics
self.modelName = "malconv"
self.experiment = args.experiment
self.log_path = args.log_path
self.save_path = args.save_path
if not os.path.isdir(self.log_path):
os.makedirs(self.log_path)
self.logger = Logger('%s/%s_%s.csv' % (self.log_path, self.modelName, args.experiment),
'epoch, time, learning_rate, tr_loss, tr_acc, val_loss, val_acc')
self.progress_bar = ProgressBar()
self.chrono = Chrono()
self.trainset, self.testset, self.trainloader, self.testloader = dataloader(args.first_n_byte)
print('==> Building model..')
self.model = MalConv(input_length=args.first_n_byte)
if torch.cuda.is_available():
self.model = torch.nn.DataParallel(self.model)
torch.backends.cudnn.benchmark = True
self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.lr)
if args.resume or self.test_only or self.dump_statistics:
self.load()
self.criterion = torch.nn.BCELoss()
self.criterion = get_torch_vars(self.criterion, False)
self.model = get_torch_vars(self.model, False)
if __name__ == '__main__':
parser = argparse.ArgumentParser(description="Train Autoencoder")
parser.add_argument("--valid",
action="store_true",
default=False,
help="Perform validation only.")
args = parser.parse_args()
chrono = Chrono()
progress_bar = ProgressBar()
_, _, trainloader, testloader = dataloader()
autoencoder = get_torch_vars(AutoEncoder(), False)
if torch.cuda.is_available():
autoencoder = torch.nn.DataParallel(autoencoder)
torch.backends.cudnn.benchmark = True
if args.valid:
valid()
exit(0)
criterion = nn.BCEWithLogitsLoss()
optimizer = optim.Adam(autoencoder.parameters())
train()
print('Finished Training')