transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, ),
(1.0, ))
]))
bs = 16
val_dataloader = torch.utils.data.DataLoader(mnist_test,
batch_size=bs,
drop_last=True,
shuffle=False)
writer = SummaryWriter('runs/' + str(args.writer))
# MODEL
if (args.model == 'Q'):
model = QMNIST((1, 28, 28), 64, 1)
elif (args.model == 'Q_PSD'):
model = QMNIST_PSD((1, 28, 28), 64, 1)
else:
model = None
device = args.device
# LOAD MODEL
with torch.no_grad():
model.load_state_dict(torch.load(args.weights, map_location='cpu'))
model = model.cuda('cuda:' + str(device))
# VALIDATION
# Loss function
mse = torch.nn.MSELoss()
lambda_reconstruction = torch.tensor([0.001
transforms.ToTensor(),
transforms.Normalize((0.5, ),
(1.0, ))
]))
bs = 64
train_dataloader = torch.utils.data.DataLoader(mnist,
batch_size=bs,
drop_last=True,
num_workers=8)
val_dataloader = torch.utils.data.DataLoader(mnist_test,
batch_size=100,
drop_last=True,
shuffle=False)
# MODEL
model = QMNIST((1, 28, 28), 64, args.model)
device = args.device
model = model.cuda('cuda:' + str(device))
# TensorboardX
writer = SummaryWriter('runs/' + str(args.writer))
# TRAINING PARAMS
n_epochs = 100
optimizer = torch.optim.Adam([{
'params': model.encoder.parameters()
}, {
'params': model.decoder.parameters()
}, {
'params': model.Q.parameters(),
'lr': 1e-2
}],
lr=1e-3)
def train_model(model, optimizer, epochs, train_dl, val_dl, wr, idx_inliers, device):
number_of_batches_per_epoch = len(iter(train_dataloader))
number_of_batches_per_epoch_validation = len(iter(val_dataloader))
# Loss function
mse = torch.nn.MSELoss()
lambda_reconstruction = torch.tensor([0.001]).cuda('cuda:'+str(device))
lambda_q = torch.tensor([1.0]).cuda('cuda:'+str(device))
# TRAINING PROCESS
for i in range(0, n_epochs):
# TRAINING
for batch_idx, (sample, label) in enumerate(train_dataloader):
inputs = sample.view(bs,1,28,28).float().cuda('cuda:'+str(device))
optimizer.zero_grad()
z, q, rec = model(inputs)
# compute loss function
reconstruction_loss = lambda_reconstruction * mse(inputs, rec)
q_loss = lambda_q * torch.sum(torch.abs(q))/bs
total_loss = torch.add(q_loss, reconstruction_loss)
# Write results
step = ((i*number_of_batches_per_epoch) + batch_idx)
norm_of_z = torch.trace(torch.matmul(z,z.t()))
# Q_0_0
# write_train_results(step, reconstruction_loss, q_loss, model.Q.get_norm_of_B(), norm_of_z, writer)
# Q_1_0
write_train_results(step, reconstruction_loss, q_loss, model.Q.get_norm_of_ATA(), norm_of_z, writer)
# Backpropagate
total_loss.backward()
optimizer.step()
if(i==2):
freeze_ENC_DEC(model)
# torch.save(model.state_dict(), os.path.join('/data/Ponc/Q_0_1/'+str(i)))
# VALIDATION
with torch.no_grad():
model = model.eval()
for batch_idx, (sample, label) in enumerate(val_dataloader):
# Separate between inliers and outliers
inputs = sample.view(100,1,28,28).float().cuda('cuda:'+str(device))
z, q, rec = model(inputs)
# compute loss function
inliers = label == idx_inliers
outliers = label != idx_inliers
inputs_in = inputs[inliers, :, :, :]
z_in = z[inliers]
q_in = q[inliers]
rec_in = rec[inliers]
rec_loss_in = lambda_reconstruction * mse(inputs_in, rec_in)
q_loss_in = lambda_q * torch.sum(torch.abs(q_in))/q_in.size()[0]
inputs_out = inputs[outliers]
z_out = z[outliers]
q_out = q[outliers]
rec_out = rec[outliers]
rec_loss_out = lambda_reconstruction * mse(inputs_out, rec_out)
q_loss_out = lambda_q * torch.sum(torch.abs(q_out))/q_out.size()[0]
step = ((i*number_of_batches_per_epoch_validation)+batch_idx)
if(q_in.size()[0]>0):
# writer.add_image('inlier/'+str(step)+'_q_'+str(q_in[0]), inputs_in[0,0,:,:].cpu().numpy().reshape(1,28,28), step)
elif(q_out.size()[0]>0):
# writer.add_image('outlier/'+str(step)+'_q_'+str(q_out[0]), inputs_out[0].cpu().numpy().reshape(1,28,28), step)
# writer.add_scalars('val_loss/rec_loss', {'inliers_rec_loss': rec_loss_in.item(),'outliers_rec_loss': rec_loss_out.item()}, step)
writer.add_scalars('val_loss/q_loss', {'inliers_q_loss': q_loss_in.item(),'outliers_q_loss': q_loss_out.item()}, step)
if __name__ == '__main__':
# Parse arguments
parser = argparse.ArgumentParser(description='Train encoder decoder to learn moment matrix.')
parser.add_argument('--model', help="Available models:\n 1. Q (learns M_inv directly)\n 2. Q_PSD (Learns M_inv = A.T*A so M is PSD)")
parser.add_argument('--writer', help="Name of the session that will be opened by tensorboard X")
parser.add_argument('--idx_inliers', help="Digit considered as inlier.")
parser.add_argument('--device', help="cuda device")
args = parser.parse_args()
# DATASETS & DATALOADERS
mnist = torchvision.datasets.MNIST('data/MNIST', train=True, download=True, transform=transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))]))
idx_inliers = int(args.idx_inliers)
mnist = select_idx(mnist, idx_inliers)
mnist_test = torchvision.datasets.MNIST('data/MNIST', train=False, transform=transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))]))
bs = 32
train_dataloader = torch.utils.data.DataLoader(mnist, batch_size=bs, drop_last=True, num_workers=8)
val_dataloader = torch.utils.data.DataLoader(mnist_test, batch_size=100, drop_last=True, shuffle=False)
# MODEL
if(args.model == 'Q'):
model = QMNIST((1,28,28), 64, 1)
elif(args.model == 'Q_PSD'):
model = QMNIST_PSD((1,28,28), 64,1)
else:
model = None
device = args.device
model = model.cuda('cuda:'+str(device))
# TensorboardX
writer = SummaryWriter('runs/'+str(args.writer))
# TRAINING PARAMS
n_epochs = 30
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
train_model(model, optimizer, n_epochs, train_dataloader, val_dataloader, writer, idx_inliers, device)