class MLP_Regression(): def __init__(self, label, parameters): super().__init__() self.writer = SummaryWriter(comment=f"_{label}_training") self.label = label self.lr = parameters['lr'] self.hidden_units = parameters['hidden_units'] self.mode = parameters['mode'] self.batch_size = parameters['batch_size'] self.num_batches = parameters['num_batches'] self.x_shape = parameters['x_shape'] self.y_shape = parameters['y_shape'] self.save_model_path = f'{parameters["save_dir"]}/{label}_model.pt' self.best_loss = np.inf self.init_net(parameters) def init_net(self, parameters): if not os.path.exists(parameters["save_dir"]): os.makedirs(parameters["save_dir"]) model_params = { 'input_shape': self.x_shape, 'classes': self.y_shape, 'batch_size': self.batch_size, 'hidden_units': self.hidden_units, 'mode': self.mode } self.net = MLP(model_params).to(DEVICE) self.optimiser = torch.optim.Adam(self.net.parameters(), lr=self.lr) self.scheduler = torch.optim.lr_scheduler.StepLR(self.optimiser, step_size=5000, gamma=0.5) print("MLP Parameters: ") print( f'batch size: {self.batch_size}, input shape: {model_params["input_shape"]}, hidden units: {model_params["hidden_units"]}, output shape: {model_params["classes"]}, lr: {self.lr}' ) def train_step(self, train_data): self.net.train() for _, (x, y) in enumerate(train_data): x, y = x.to(DEVICE), y.to(DEVICE) self.net.zero_grad() self.loss_info = torch.nn.functional.mse_loss(self.net(x), y, reduction='sum') self.loss_info.backward() self.optimiser.step() self.epoch_loss = self.loss_info.item() def evaluate(self, x_test): self.net.eval() with torch.no_grad(): y_test = self.net(x_test.to(DEVICE)).detach().cpu().numpy() return y_test def log_progress(self, step): write_loss(self.writer, self.loss_info, step)
class Greedy_Bandit(Bandit): def __init__(self, label, *args): super().__init__(label, *args) self.writer = SummaryWriter(comment=f"_{label}_training"), def init_net(self, parameters): model_params = { 'input_shape': self.x.shape[1] + 2, 'classes': 1 if len(self.y.shape) == 1 else self.y.shape[1], 'batch_size': self.batch_size, 'hidden_units': parameters['hidden_units'], 'mode': parameters['mode'] } self.net = MLP(model_params).to(DEVICE) self.optimiser = torch.optim.Adam(self.net.parameters(), lr=self.lr) self.scheduler = torch.optim.lr_scheduler.StepLR(self.optimiser, step_size=5000, gamma=0.5) print(f'Bandit {self.label} Parameters: ') print( f'buffer_size: {self.buffer_size}, batch size: {self.batch_size}, number of samples: {self.n_samples}, epsilon: {self.epsilon}' ) print("MLP Parameters: ") print( f'input shape: {model_params["input_shape"]}, hidden units: {model_params["hidden_units"]}, output shape: {model_params["classes"]}, lr: {self.lr}' ) def loss_step(self, x, y, batch_id): self.net.train() self.net.zero_grad() net_loss = torch.nn.functional.mse_loss(self.net(x).squeeze(), y, reduction='sum') net_loss.backward() self.optimiser.step() return net_loss def log_progress(self, step): write_loss(self.writer[0], self.loss_info, step) self.writer[0].add_scalar('logs/cumulative_regret', self.cumulative_regrets[-1], step)
class MLP_Classification(): def __init__(self, label, parameters): super().__init__() self.writer = SummaryWriter(comment=f"_{label}_training") self.label = label self.lr = parameters['lr'] self.hidden_units = parameters['hidden_units'] self.mode = parameters['mode'] self.batch_size = parameters['batch_size'] self.num_batches = parameters['num_batches'] self.x_shape = parameters['x_shape'] self.classes = parameters['classes'] self.save_model_path = f'{parameters["save_dir"]}/{label}_model.pt' self.best_acc = 0. self.dropout = parameters['dropout'] self.init_net(parameters) def init_net(self, parameters): if not os.path.exists(parameters["save_dir"]): os.makedirs(parameters["save_dir"]) model_params = { 'input_shape': self.x_shape, 'classes': self.classes, 'batch_size': self.batch_size, 'hidden_units': self.hidden_units, 'mode': self.mode, 'dropout': self.dropout, } if self.dropout: self.net = MLP_Dropout(model_params).to(DEVICE) print('MLP Dropout Parameters: ') print(f'batch size: {self.batch_size}, input shape: {model_params["input_shape"]}, hidden units: {model_params["hidden_units"]}, output shape: {model_params["classes"]}, lr: {self.lr}') else: self.net = MLP(model_params).to(DEVICE) print('MLP Parameters: ') print(f'batch size: {self.batch_size}, input shape: {model_params["input_shape"]}, hidden units: {model_params["hidden_units"]}, output shape: {model_params["classes"]}, lr: {self.lr}') self.optimiser = torch.optim.SGD(self.net.parameters(), lr=self.lr) self.scheduler = torch.optim.lr_scheduler.StepLR(self.optimiser, step_size=100, gamma=0.5) def train_step(self, train_data): self.net.train() for _, (x, y) in enumerate(tqdm(train_data)): x, y = x.to(DEVICE), y.to(DEVICE) self.net.zero_grad() self.loss_info = torch.nn.functional.cross_entropy(self.net(x), y, reduction='sum') self.loss_info.backward() self.optimiser.step() def predict(self, X): probs = torch.nn.Softmax(dim=1)(self.net(X)) preds = torch.argmax(probs, dim=1) return preds, probs def evaluate(self, test_loader): self.net.eval() print('Evaluating on validation data') correct = 0 total = 0 with torch.no_grad(): for data in tqdm(test_loader): X, y = data X, y = X.to(DEVICE), y.to(DEVICE) preds, _ = self.predict(X) total += self.batch_size correct += (preds == y).sum().item() self.acc = correct / total print(f'{self.label} validation accuracy: {self.acc}') def log_progress(self, step): write_loss(self.writer, self.loss_info, step) write_acc(self.writer, self.acc, step)
class VanillaAE(nn.Module): def __init__(self, opt): super(VanillaAE, self).__init__() self.opt = opt self.device = torch.device("cuda:0" if not opt.no_cuda else "cpu") nc = int(opt.nc) imageSize = int(opt.imageSize) nz = int(opt.nz) nblk = int(opt.nblk) # generator self.netG = MLP(input_dim=nc * imageSize * imageSize, output_dim=nc * imageSize * imageSize, dim=nz, n_blk=nblk, norm='none', activ='relu').to(self.device) weights_init(self.netG) if opt.netG != '': self.netG.load_state_dict( torch.load(opt.netG, map_location=self.device)) print_and_write_log(opt.train_log_file, 'netG:') print_and_write_log(opt.train_log_file, str(self.netG)) # losses self.criterion = nn.MSELoss() # define focal frequency loss self.criterion_freq = FFL(loss_weight=opt.ffl_w, alpha=opt.alpha, patch_factor=opt.patch_factor, ave_spectrum=opt.ave_spectrum, log_matrix=opt.log_matrix, batch_matrix=opt.batch_matrix).to( self.device) # misc self.to(self.device) # optimizer self.optimizerG = optim.Adam(self.netG.parameters(), lr=opt.lr, betas=(opt.beta1, opt.beta2)) def forward(self): pass def gen_update(self, data, epoch, matrix=None): self.netG.zero_grad() real = data.to(self.device) if matrix is not None: matrix = matrix.to(self.device) recon = self.netG(real) # apply pixel-level loss errG_pix = self.criterion(recon, real) * self.opt.mse_w # apply focal frequency loss if epoch >= self.opt.freq_start_epoch: errG_freq = self.criterion_freq(recon, real, matrix) else: errG_freq = torch.tensor(0.0).to(self.device) errG = errG_pix + errG_freq errG.backward() self.optimizerG.step() return errG_pix, errG_freq def sample(self, x): x = x.to(self.device) self.netG.eval() with torch.no_grad(): recon = self.netG(x) self.netG.train() return recon def save_checkpoints(self, ckpt_dir, epoch): torch.save(self.netG.state_dict(), '%s/netG_epoch_%03d.pth' % (ckpt_dir, epoch))