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))