class VanillaClassifier(): def __init__(self, config): # training device - try to find a gpu, if not just use cpu self.device = torch.device( 'cuda:0' if torch.cuda.is_available() else 'cpu') print('[INFO]: using \'{}\' device'.format(self.device)) # initialize model self.model = Classifier(config['input_dimensions'], config['output_dimension'], hid_act=config['hidden_activation'], norm=config['normalization']) # if model file provided, load pretrained params if config['model_file']: self.load(config['model_file']) # move the model to the training device self.model.to(self.device) self.config = config def load(self, model_file): self.model.load_state_dict( torch.load(model_file, map_location=self.device)) print('[INFO]: loaded model from \'{}\''.format(model_file)) def logits(self, x): return self.model(x) def predict(self, x): return torch.argmax(self.model(x), dim=1) def train_epochs(self, train_loader, test_loader): # define cross entropy loss (requires logits as outputs) loss_fn = torch.nn.CrossEntropyLoss() # initialize an optimizer optimizer = torch.optim.Adam(self.model.parameters(), lr=self.config['learning_rate'], weight_decay=self.config['weight_decay']) # initialize tensorboard writer writer = SummaryWriter('{}runs/{}/'.format( self.config['output_directory'], self.config['model_name'])) print('[INFO]: training...') # train through all epochs for e in range(self.config['number_epochs']): # get epoch start time epoch_start = time.time() # reset accumulators train_epoch_loss = 0.0 train_num_correct = 0 test_epoch_loss = 0.0 test_num_correct = 0 # run through epoch of train data for i, batch in enumerate(train_loader): # parse batch and move to training device input_batch = batch['image'].to(self.device) label_batch = batch['label'].to(self.device) # compute output batch logits and predictions logits_batch = self.model(input_batch) pred_batch = torch.argmax(logits_batch, dim=1) # compute loss loss = loss_fn(logits_batch, label_batch) # accumulate loss train_epoch_loss += loss.item() # accumulate number correct train_num_correct += torch.sum( torch.tensor(pred_batch == label_batch)).item() # zero out gradient attributes for all trainable params optimizer.zero_grad() # compute gradients w.r.t loss (repopulate gradient # attribute for all trainable params) loss.backward() # update params with current gradients optimizer.step() # compute epoch average loss and accuracy metrics train_loss = train_epoch_loss / i train_acc = 100.0 * train_num_correct / self.config['number_train'] # run through epoch of test data for i, batch in enumerate(test_loader): # parse batch and move to training device input_batch = batch['image'].to(self.device) label_batch = batch['label'].to(self.device) # compute output batch logits and predictions logits_batch = self.model(input_batch) pred_batch = torch.argmax(logits_batch, dim=1) # compute loss loss = loss_fn(logits_batch, label_batch) # accumulate loss test_epoch_loss += loss.item() # accumulate number correct test_num_correct += torch.sum( torch.tensor(pred_batch == label_batch)).item() # compute epoch average loss and accuracy metrics test_loss = test_epoch_loss / i test_acc = 100.0 * test_num_correct / self.config['number_test'] # compute epoch time epoch_time = time.time() - epoch_start # save model torch.save( self.model.state_dict(), '{}{}.pt'.format(self.config['output_directory'], self.config['model_name'])) # add metrics to tensorboard writer.add_scalar('Loss/Train', train_loss, e + 1) writer.add_scalar('Accuracy/Train', train_acc, e + 1) writer.add_scalar('Loss/Test', test_loss, e + 1) writer.add_scalar('Accuracy/Test', test_acc, e + 1) # print epoch metrics template = '[INFO]: Epoch {}, Epoch Time {:.2f}s, ' \ 'Train Loss: {:.2f}, Train Accuracy: {:.2f}, ' \ 'Test Loss: {:.2f}, Test Accuracy: {:.2f}' print( template.format(e + 1, epoch_time, train_loss, train_acc, test_loss, test_acc))
class AdversarialClassifier(): def __init__(self, config): # training device - try to find a gpu, if not just use cpu self.device = torch.device( 'cuda:0' if torch.cuda.is_available() else 'cpu') print('[INFO]: using \'{}\' device'.format(self.device)) # initialize model self.model = Classifier( config['input_dimensions'], config['output_dimension'], hid_act=config['hidden_activation'], norm=config['normalization']) # if model file provided, load pretrained params if config['model_file']: self.load(config['model_file']) # move the model to the training device self.model.to(self.device) self.config = config def load(self, model_file): self.model.load_state_dict( torch.load(model_file, map_location=self.device)) print('[INFO]: loaded model from \'{}\'' .format(model_file)) def logits(self, x): return self.model(x) def predict(self, x): return torch.argmax(self.model(x), dim=1) def train_epochs(self, train_loader, test_loader): # define cross entropy loss (requires logits as outputs) loss_fn = torch.nn.CrossEntropyLoss() # initialize an optimizer optimizer = torch.optim.Adam( self.model.parameters(), lr=self.config['learning_rate'], weight_decay=self.config['weight_decay']) # initialize guide model (for making adversarial samples) guide_model = Classifier( self.config['input_dimensions'], self.config['output_dimension'], hid_act=self.config['hidden_activation'], norm=self.config['normalization']) # must be able to load guide model to proceed guide_model.load_state_dict( torch.load(self.config['guide_model_file'], map_location=self.device)) print('[INFO]: loaded guide model from \'{}\'' .format(self.config['guide_model_file'])) # move guide model to the training device guide_model.to(self.device) # initialize tensorboard writer writer = SummaryWriter('{}runs/{}/'.format( self.config['output_directory'], self.config['model_name'])) print('[INFO]: training...') # train through all epochs for e in range(self.config['number_epochs']): # get epoch start time epoch_start = time.time() # reset accumulators train_epoch_loss = 0.0 train_num_correct = 0 test_epoch_loss = 0.0 test_num_correct = 0 # run through epoch of train data for i, batch in enumerate(train_loader): # parse batch and move to training device input_batch = batch['image'].to(self.device) label_batch = batch['label'].to(self.device) # require gradient for input data (need to do this to compute # the gradients for inputs during backward() call) input_batch.requires_grad = True # make adversarial samples from input batch adv_logits_batch = guide_model(input_batch) adv_loss = loss_fn(adv_logits_batch, label_batch) adv_loss.backward() adv_grads = input_batch.grad epsilon = (0.5 * ((2 * torch.rand(1)) - 1)).to(self.device) adv_input_batch = input_batch + ( epsilon * torch.sign(adv_grads)) # keep pixel values of adv batch in [-1, 1] new_min, new_max = -1., 1. old_min = torch.min(adv_input_batch) old_max = torch.max(adv_input_batch) adv_input_batch = (((adv_input_batch - old_min) / ( old_max - old_min)) * (new_max - new_min)) + new_min # compute output batch logits and predictions logits_batch = self.model(input_batch) adv_logits_batch = self.model(adv_input_batch) pred_batch = torch.argmax(logits_batch, dim=1) # compute combined normal/adversarial loss loss = loss_fn(logits_batch, label_batch) adv_loss = loss_fn(adv_logits_batch, label_batch) alpha = 0.5 total_loss = (alpha * loss) + ((1 - alpha) * adv_loss) # accumulate loss train_epoch_loss += loss.item() # accumulate number correct train_num_correct += torch.sum( torch.tensor(pred_batch == label_batch)).item() # zero out gradient attributes for all trainable params optimizer.zero_grad() # compute gradients w.r.t loss (repopulate gradient # attribute for all trainable params) total_loss.backward() # update params with current gradients optimizer.step() # compute epoch average loss and accuracy metrics train_loss = train_epoch_loss / i train_acc = 100.0 * train_num_correct / self.config['number_train'] # run through epoch of test data for i, batch in enumerate(test_loader): # parse batch and move to training device input_batch = batch['image'].to(self.device) label_batch = batch['label'].to(self.device) # compute output batch logits and predictions logits_batch = self.model(input_batch) pred_batch = torch.argmax(logits_batch, dim=1) # compute loss loss = loss_fn(logits_batch, label_batch) # accumulate loss test_epoch_loss += loss.item() # accumulate number correct test_num_correct += torch.sum( torch.tensor(pred_batch == label_batch)).item() # compute epoch average loss and accuracy metrics test_loss = test_epoch_loss / i test_acc = 100.0 * test_num_correct / self.config['number_test'] # compute epoch time epoch_time = time.time() - epoch_start # save model torch.save(self.model.state_dict(), '{}{}.pt'.format( self.config['output_directory'], self.config['model_name'])) # add metrics to tensorboard writer.add_scalar('Loss/Train', train_loss, e + 1) writer.add_scalar('Accuracy/Train', train_acc, e + 1) writer.add_scalar('Loss/Test', test_loss, e + 1) writer.add_scalar('Accuracy/Test', test_acc, e + 1) # print epoch metrics template = '[INFO]: Epoch {}, Epoch Time {:.2f}s, ' \ 'Train Loss: {:.2f}, Train Accuracy: {:.2f}, ' \ 'Test Loss: {:.2f}, Test Accuracy: {:.2f}' print(template.format(e + 1, epoch_time, train_loss, train_acc, test_loss, test_acc))
class SmoothClassifier(): def __init__(self, config): # training device - try to find a gpu, if not just use cpu self.device = torch.device( 'cuda:0' if torch.cuda.is_available() else 'cpu') print('[INFO]: using \'{}\' device'.format(self.device)) # initialize model self.model = Classifier( config['input_dimensions'], config['output_dimension'], hid_act=config['hidden_activation'], norm=config['normalization']) # initialize a random input distribution self.n_samples = 32 mean = 0.0 stdev = 0.25 self.input_dist = torch.distributions.Normal( mean * torch.ones( config['batch_size'], self.n_samples, config['input_dimensions'][-1], config['input_dimensions'][0], config['input_dimensions'][1]), stdev * torch.ones( config['batch_size'], self.n_samples, config['input_dimensions'][-1], config['input_dimensions'][0], config['input_dimensions'][1]) ) # if model file provided, load pretrained params if config['model_file']: self.load(config['model_file']) # move the model to the training device self.model.to(self.device) self.config = config def load(self, model_file): self.model.load_state_dict( torch.load(model_file, map_location=self.device)) print('[INFO]: loaded model from \'{}\'' .format(model_file)) def logits(self, x): return self.model(x) def predict(self, x): # get number of samples in batch bs = x.shape[0] # add noise to input batch x = x.unsqueeze(1) + self.input_dist.sample()[:bs].to(self.device) # reshape input batch by stacking samples into batch dimension x = x.view(( bs * self.n_samples, self.config['input_dimensions'][-1], self.config['input_dimensions'][0], self.config['input_dimensions'][1])) # compute output batch logits and predictions logits = self.model(x) pred = torch.argmax(logits, dim=1) # reshape predictions to unstack samples from batch dimension pred = pred.view((bs, self.n_samples)) # take mode along sample dim to get final prediction pred = torch.mode(pred, dim=1)[0] return pred def train_epochs(self, train_loader, test_loader): # define cross entropy loss (requires logits as outputs) loss_fn = torch.nn.CrossEntropyLoss() # initialize an optimizer optimizer = torch.optim.Adam( self.model.parameters(), lr=self.config['learning_rate'], weight_decay=self.config['weight_decay']) # initialize tensorboard writer writer = SummaryWriter('{}runs/{}/'.format( self.config['output_directory'], self.config['model_name'])) print('[INFO]: training...') # train through all epochs for e in range(self.config['number_epochs']): # get epoch start time epoch_start = time.time() # reset accumulators train_epoch_loss = 0.0 train_num_correct = 0 test_epoch_loss = 0.0 test_num_correct = 0 # run through epoch of train data for i, batch in enumerate(train_loader): # parse batch and move to training device input_batch = batch['image'].to(self.device) label_batch = batch['label'].to(self.device) # get number of samples in batch bs = input_batch.shape[0] # add noise to input batch input_batch = input_batch.unsqueeze(1) \ + self.input_dist.sample()[:bs].to(self.device) # reshape input batch by stacking samples into batch dimension input_batch = input_batch.view(( bs * self.n_samples, self.config['input_dimensions'][-1], self.config['input_dimensions'][0], self.config['input_dimensions'][1])) # keep pixel values of batch in [-1, 1] new_min, new_max = -1., 1. old_min = torch.min(input_batch) old_max = torch.max(input_batch) input_batch = (((input_batch - old_min) / ( old_max - old_min)) * (new_max - new_min)) + new_min # repeat and interleave label batch to repeat labels for each # samples stacked into batch dimension label_batch = label_batch.repeat_interleave(self.n_samples) # compute output batch logits and predictions logits_batch = self.model(input_batch) pred_batch = torch.argmax(logits_batch, dim=1) # compute loss loss = loss_fn(logits_batch, label_batch) # accumulate loss train_epoch_loss += loss.item() # accumulate number correct train_num_correct += torch.sum( torch.tensor(pred_batch == label_batch)).item() # zero out gradient attributes for all trainable params optimizer.zero_grad() # compute gradients w.r.t loss (repopulate gradient # attribute for all trainable params) loss.backward() # update params with current gradients optimizer.step() # compute epoch average loss and accuracy metrics train_loss = train_epoch_loss / i train_acc = 100.0 * train_num_correct \ / (self.config['number_train'] * self.n_samples) # run through epoch of test data for i, batch in enumerate(test_loader): # parse batch and move to training device input_batch = batch['image'].to(self.device) label_batch = batch['label'].to(self.device) # compute output batch logits and predictions logits_batch = self.model(input_batch) pred_batch = torch.argmax(logits_batch, dim=1) # compute loss loss = loss_fn(logits_batch, label_batch) # accumulate loss test_epoch_loss += loss.item() # accumulate number correct test_num_correct += torch.sum( torch.tensor(pred_batch == label_batch)).item() # compute epoch average loss and accuracy metrics test_loss = test_epoch_loss / i test_acc = 100.0 * test_num_correct / self.config['number_test'] # compute epoch time epoch_time = time.time() - epoch_start # save model torch.save(self.model.state_dict(), '{}{}.pt'.format( self.config['output_directory'], self.config['model_name'])) # add metrics to tensorboard writer.add_scalar('Loss/Train', train_loss, e + 1) writer.add_scalar('Accuracy/Train', train_acc, e + 1) writer.add_scalar('Loss/Test', test_loss, e + 1) writer.add_scalar('Accuracy/Test', test_acc, e + 1) # print epoch metrics template = '[INFO]: Epoch {}, Epoch Time {:.2f}s, ' \ 'Train Loss: {:.2f}, Train Accuracy: {:.2f}, ' \ 'Test Loss: {:.2f}, Test Accuracy: {:.2f}' print(template.format(e + 1, epoch_time, train_loss, train_acc, test_loss, test_acc))