def test(model, test_loader):
"""
Evaluate model on the test set
Args:
model: pytorch MDN model
test_loader: loader for the test dataset
Returns:
accuracy on the test dataset
"""
error_sum = 0
for batch_idx, (minibatch, labels) in enumerate(test_loader):
minibatch = minibatch.reshape(batch_size, 1, 784)
labels = labels.reshape(batch_size,1)
labels = labels.int()
pi, sigma, mu = model(minibatch)
samples = mdn.sample(pi, sigma, mu).int()
error = (samples != labels)
error_sum = error_sum + error.sum()
return 1 - error_sum.item()/(len(test_loader)*test_loader.batch_size)
def evaluate(self, save_dir='data/', prefix=''):
self.load() # load the model as constructed
cuda = True if torch.cuda.is_available() else False
if cuda:
self.model.cuda()
self.model.eval()
saved_model_str = self.saved_model.replace('/', '_') + prefix
# Get the file names
Ypred_file = os.path.join(save_dir, 'test_Ypred_{}.csv'.format(saved_model_str))
Xtruth_file = os.path.join(save_dir, 'test_Xtruth_{}.csv'.format(saved_model_str))
Ytruth_file = os.path.join(save_dir, 'test_Ytruth_{}.csv'.format(saved_model_str))
Xpred_file = os.path.join(save_dir, 'test_Xpred_{}.csv'.format(saved_model_str))
# keep time
tk = time_keeper(os.path.join(save_dir, 'evaluation_time.txt'))
# Open those files to append
with open(Xtruth_file, 'a') as fxt,open(Ytruth_file, 'a') as fyt,\
open(Ypred_file, 'a') as fyp, open(Xpred_file, 'a') as fxp:
# Loop through the eval data and evaluate
for ind, (geometry, spectra) in enumerate(self.test_loader):
if cuda:
geometry = geometry.cuda()
spectra = spectra.cuda()
# Initialize the geometry first
print('model in eval:', self.model)
pi, sigma, mu = self.model(spectra) # Get the output
Xpred = mdn.sample(pi, sigma, mu).detach().cpu().numpy()
tk.record(1)
return Ypred_file, Ytruth_file
def predict(self, Ytruth_file, save_dir='data/', prefix=''):
self.load() # load the model as constructed
cuda = True if torch.cuda.is_available() else False
if cuda:
self.model.cuda()
self.model.eval()
saved_model_str = self.saved_model.replace('/', '_') + prefix
Ytruth = pd.read_csv(Ytruth_file, header=None, delimiter=',') # Read the input
if len(Ytruth.columns) == 1: # The file is not delimitered by ',' but ' '
Ytruth = pd.read_csv(Ytruth_file, header=None, delimiter=' ')
Ytruth_tensor = torch.from_numpy(Ytruth.values).to(torch.float)
print('shape of Ytruth tensor :', Ytruth_tensor.shape)
# Get the file names
Ypred_file = os.path.join(save_dir, 'test_Ypred_{}.csv'.format(saved_model_str))
Ytruth_file = os.path.join(save_dir, 'test_Ytruth_{}.csv'.format(saved_model_str))
Xpred_file = os.path.join(save_dir, 'test_Xpred_{}.csv'.format(saved_model_str))
# keep time
tk = time_keeper(os.path.join(save_dir, 'evaluation_time.txt'))
if cuda:
Ytruth_tensor = Ytruth_tensor.cuda()
print('model in eval:', self.model)
pi, sigma, mu = self.model(Ytruth_tensor) # Get the output
Xpred = mdn.sample(pi, sigma, mu).detach().cpu().numpy()
# Open those files to append
with open(Ytruth_file, 'a') as fyt, open(Ypred_file, 'a') as fyp, open(Xpred_file, 'a') as fxp:
np.savetxt(fyt, Ytruth_tensor.cpu().data.numpy())
np.savetxt(fxp, Xpred)
if self.flags.data_set != 'Yang_sim':
Ypred = simulator(self.flags.data_set, Xpred)
np.savetxt(fyp, Ypred)
tk.record(1)
return Ypred_file, Ytruth_file
def train(self):
"""
The major training function. This would start the training using information given in the flags
:return: None
"""
cuda = True if torch.cuda.is_available() else False
if cuda:
self.model.cuda()
# Construct optimizer after the model moved to GPU
self.optm = self.make_optimizer()
self.lr_scheduler = self.make_lr_scheduler(self.optm)
# Time keeping
tk = time_keeper(
time_keeping_file=os.path.join(self.ckpt_dir, 'training time.txt'))
for epoch in range(self.flags.train_step):
# Set to Training Mode
train_loss = 0
# boundary_loss = 0 # Unnecessary during training since we provide geometries
self.model.train()
for j, (geometry, spectra) in enumerate(self.train_loader):
if cuda:
geometry = geometry.cuda() # Put data onto GPU
spectra = spectra.cuda() # Put data onto GPU
#print('spectra = ', spectra)
#print('geometry = ', geometry)
self.optm.zero_grad() # Zero the gradient first
pi, sigma, mu = self.model(spectra) # Get the output
#print('spectra = {}, pi, sigma, mu = {}, {}, {}'.format(spectra.cpu().numpy(),
# pi.detach().cpu().numpy()[0,:],
# sigma.detach().cpu().numpy()[0,:,0],
# mu.detach().cpu().numpy()[0,:,0]))
#print('geometry shape', geometry.size())
loss = self.make_loss(pi, sigma, mu,
geometry) # Get the loss tensor
#loss = self.make_loss(pi, sigma, mu, geometry, warmup=epoch) # Get the loss tensor
#Xpred = mdn.sample(pi, sigma, mu).detach().cpu().numpy()
#Ypred = torch.tensor(simulator(self.flags.data_set, Xpred), requires_grad=False)
#if cuda:
# Ypred = Ypred.cuda()
#simulator_loss = nn.functional.mse_loss(Ypred, spectra).detach().cpu().numpy() # Get the loss tensor
#print('nll loss at epoch {}, batch {} is {} '.format(epoch, j, loss.detach().cpu().numpy()))
loss.backward() # Calculate the backward gradients
# gradient clipping
torch.nn.utils.clip_grad_value_(self.model.parameters(), 1)
self.optm.step() # Move one step the optimizer
train_loss += loss # Aggregate the loss
# boundary_loss += self.Boundary_loss # Aggregate the BDY loss
# Calculate the avg loss of training
train_avg_loss = train_loss.cpu().data.numpy() / (j + 1)
# boundary_avg_loss = boundary_loss.cpu().data.numpy() / (j + 1)
if epoch % self.flags.eval_step == 0: # For eval steps, do the evaluations and tensor board
# Record the training loss to the tensorboard
self.log.add_scalar('Loss/train', train_avg_loss, epoch)
#self.log.add_scalar('Loss/simulator_train', simulator_loss, epoch)
# self.log.add_scalar('Loss/BDY_train', boundary_avg_loss, epoch)
# Set to Evaluation Mode
self.model.eval()
print("Doing Evaluation on the model now")
test_loss = 0
for j, (geometry, spectra) in enumerate(
self.test_loader): # Loop through the eval set
if cuda:
geometry = geometry.cuda()
spectra = spectra.cuda()
pi, sigma, mu = self.model(spectra) # Get the output
if self.flags.data_set == 'meta_material':
loss = self.make_loss(pi, sigma, mu,
geometry) # Get the loss tensor
test_loss += loss.detach().cpu().numpy()
else:
Xpred = mdn.sample(pi, sigma, mu).numpy()
Ypred_np = simulator(self.flags.data_set, Xpred)
mae, mse = compare_truth_pred(Ypred_np,
spectra.cpu().numpy(),
cut_off_outlier_thres=10,
quiet_mode=True)
test_loss += np.mean(mse) # Aggregate the loss
break
# only get the first batch that is enough
# Record the testing loss to the tensorboard
test_avg_loss = test_loss / (j + 1)
self.log.add_scalar('Loss/test', test_avg_loss, epoch)
print("This is Epoch %d, training loss %.5f, validation loss %.5f"\
% (epoch, train_avg_loss, test_avg_loss ))
#print("This is Epoch %d, training loss %.5f, validation loss %.5f, training simulator loss %.5f" \
# % (epoch, train_avg_loss, test_avg_loss, simulator_loss ))
# Plotting the first spectra prediction for validation
# f = self.compare_spectra(Ypred=logit[0,:].cpu().data.numpy(), Ytruth=spectra[0,:].cpu().data.numpy())
# self.log.add_figure(tag='spectra compare',figure=f,global_step=epoch)
# Model improving, save the model down
if test_avg_loss < self.best_validation_loss:
self.best_validation_loss = test_avg_loss
self.save()
print("Saving the model down...")
if self.best_validation_loss < self.flags.stop_threshold:
print("Training finished EARLIER at epoch %d, reaching loss of %.5f" %\
(epoch, self.best_validation_loss))
return None
# Learning rate decay upon plateau
self.lr_scheduler.step(train_avg_loss)
self.log.close()
tk.record(1) # Record the total time of the training peroid
for epoch in range(num_epochs):
for batch_idx, (labels, minibatch) in enumerate(train_loader):
gt = labels[0:3]
minibatch = [[numb == minibatch[i] for numb in range(10)]
for i in range(batch_size)]
minibatch = torch.FloatTensor(minibatch).unsqueeze(1)
labels = labels.reshape(batch_size, 784)
model.zero_grad()
pi, sigma, mu = model(minibatch)
loss = mdn.mdn_loss(pi, sigma, mu, labels)
loss.backward()
optimizer.step()
print('Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, batch_idx * len(minibatch), len(train_loader.dataset),
100. * batch_idx / len(train_loader), loss.item()))
# Visualize
if batch_idx == 0:
samples = mdn.sample(pi, sigma, mu).int()
images = samples[0:3]
plot(gt, images, epoch)
print(
"The training is compelete! \nVisualization results have been saved in the folder 'figures'"
)
alp0, mu0, sigma0 = model0(minibatch)
loss0 = mdn.loss_NLL_MDN(alp0, mu0, sigma0, label)
# mse_loss = nn.MSELoss()
# loss0 = mse_loss(model0(minibatch), labels.float())
loss0.backward()
optimizer0.step()
print('\n')
# sample new points from the trained model
t_test = ts(np.random.sample(int(5e3))[:, np.newaxis].astype(np.float32))
t_test0 = ts(np.random.sample(int(2e3))[:, np.newaxis].astype(np.float32))
alp, mu, sigma = model(t_test)
samples = mdn.sample(alp, mu, sigma).detach().numpy()
alp0, mu0, sigma0 = model0(t_test0)
samples0 = mdn.sample(alp0, mu0, sigma0).detach().numpy()
# samples0 = model0(t_test0).detach().numpy()
plt.scatter(x_data, y_data, marker='.')
plt.scatter(samples0[:,0], samples0[:,1], marker='.')
plt.scatter(samples[:,0], samples[:,1], marker='.')
plt.scatter(np.concatenate((x_mean1,x_mean2),axis=0),
np.concatenate((y_mean1,y_mean2),axis=0), marker='x')
plt.arrow(6,0, 0,5, head_width=0.05, head_length=0.1)
plt.xlabel("x")
#x = torch.rand(n_batch, batch_size, in_feat)
#y = torch.rand(n_batch, batch_size, out_feat)
# train the model
for epoch in range(0, n_epoch):
e_loss = 0.0
for minibatch, labels in zip(X_train, Y_train):
minibatch = Variable(minibatch).float()
labels = Variable(labels).float()
model.zero_grad()
pi, sigma, mu = model(minibatch)
loss = mdn.mdn_loss(pi, sigma, mu, labels)
loss.backward()
optimizer.step()
e_loss += loss.data[0]
print("epoch: ", epoch, ": ", e_loss)
# sampleing new points from the trained model
# first predict parameters from MDN
pi, sigma, mu = model(Variable(X_test[0]).float())
#sample 100 instances using parameters
samples_0 = mdn.sample(pi, sigma, mu)
print("samples: ", samples_0.size())
print("done!")
