def train(self, data_loader, epochs):
"""
Trains Neural Process.
Parameters
----------
dataloader : torch.utils.DataLoader instance
epochs : int
Number of epochs to train for.
"""
for epoch in range(epochs):
epoch_loss = 0.
for i, data in enumerate(data_loader):
self.optimizer.zero_grad()
# Sample number of context and target points
num_context = randint(*self.num_context_range)
num_extra_target = randint(*self.num_extra_target_range)
# Create context and target points and apply neural process
if self.is_img:
img, _ = data # data is a tuple (img, label)
batch_size = img.size(0)
context_mask, target_mask = \
batch_context_target_mask(self.neural_process.img_size,
num_context, num_extra_target,
batch_size)
img = img.to(self.device)
context_mask = context_mask.to(self.device)
target_mask = target_mask.to(self.device)
p_y_pred, q_target, q_context = \
self.neural_process(img, context_mask, target_mask)
# Calculate y_target as this will be required for loss
_, y_target = img_mask_to_np_input(img, target_mask)
else:
x, y = data
x_context, y_context, x_target, y_target = \
context_target_split(x, y, num_context, num_extra_target)
p_y_pred, q_target, q_context = \
self.neural_process(x_context, y_context, x_target, y_target)
loss = self._loss(p_y_pred, y_target, q_target, q_context)
loss.backward()
self.optimizer.step()
epoch_loss += loss.item()
self.steps += 1
if self.steps % self.print_freq == 0:
print("iteration {}, loss {:.3f}".format(
self.steps, loss.item()))
print("Epoch: {}, Avg_loss: {}".format(
epoch, epoch_loss / len(data_loader)))
self.epoch_loss_history.append(epoch_loss / len(data_loader))
def forward(self, img, context_mask, target_mask):
"""
Given an image and masks of context and target points, returns a
distribution over pixel intensities at the target points.
Parameters
----------
img : torch.Tensor
Shape (batch_size, channels, height, width)
context_mask : torch.ByteTensor
Shape (batch_size, height, width). Binary mask indicating
the pixels to be used as context.
target_mask : torch.ByteTensor
Shape (batch_size, height, width). Binary mask indicating
the pixels to be used as target.
"""
x_context, y_context = img_mask_to_np_input(img, context_mask)
x_target, y_target = img_mask_to_np_input(img, target_mask)
return self.neural_process(x_context, y_context, x_target, y_target)
def train(self, data_loader, epochs, x_context_plot, y_context_plot):
"""
Trains Neural Process.
Parameters
----------
dataloader : torch.utils.DataLoader instance
epochs : int
Number of epochs to train for.
"""
for epoch in range(epochs):
epoch_loss = 0.
for i, data in enumerate(data_loader):
self.optimizer.zero_grad()
# Sample number of context and target points
num_context = randint(*self.num_context_range)
num_extra_target = randint(*self.num_extra_target_range)
# Create context and target points and apply neural process
if self.is_img:
img, _ = data # data is a tuple (img, label)
batch_size = img.size(0)
context_mask, target_mask = \
batch_context_target_mask(self.neural_process.img_size,
num_context, num_extra_target,
batch_size)
img = img.to(self.device)
context_mask = context_mask.to(self.device)
target_mask = target_mask.to(self.device)
p_y_pred, q_target, q_context = \
self.neural_process(img, context_mask, target_mask)
# Calculate y_target as this will be required for loss
_, y_target = img_mask_to_np_input(img, target_mask)
else:
x, y = data
x_context, y_context, x_target, y_target = \
context_target_split(x, y, num_context, num_extra_target)
p_y_pred, q_target, q_context = \
self.neural_process(x_context, y_context, x_target, y_target)
loss = self._loss(p_y_pred, y_target, q_target, q_context, MMD)
loss.backward()
self.optimizer.step()
epoch_loss += loss.item()
self.steps += 1
if self.steps % self.print_freq == 0:
print("iteration {}, loss {:.3f}".format(
self.steps, loss.item()))
print("Epoch: {}, Avg_loss: {}".format(
epoch, epoch_loss / len(data_loader)))
self.epoch_loss_history.append(epoch_loss / len(data_loader))
self.neural_process.training = False
for i in range(64):
# Neural process returns distribution over y_target
p_y_pred = self.neural_process(x_context_plot, y_context_plot,
self.x_target_plot)
# Extract mean of distribution
mu = p_y_pred.loc.detach()
std = p_y_pred.stddev.detach()
plt.plot(self.x_target_plot.numpy()[0],
mu.numpy()[0],
alpha=0.05,
c='b')
plt.scatter(x_context_plot[0].numpy(),
y_context_plot[0].numpy(),
c='k')
if self.MMD == True:
st = 'MMD'
else:
st = 'KLD'
if self.fixed_sigma == True:
ss = 'FIXED_SIGMA={}'.format(self.sig)
else:
ss = None
plt.savefig('./NewNewPlots/{}{}alpha{}epoch{}.png'.format(
ss, st, self.alpha, epoch))
plt.clf()
plt.plot(self.x_target_plot.numpy()[0], std.numpy()[0], c='b')
plt.savefig('./NewNewPlots/{}{}alpha{}epoch{}_variance.png'.format(
ss, st, self.alpha, epoch))
plt.clf()
self.neural_process.training = True
