def valid_step(model, dataset, device, criterion, batch_size):
model.eval()
y_preds = []
losses = []
batch_sizes = []
tic = time.time()
with torch.no_grad():
for Xi, yi in torch.utils.data.DataLoader(
dataset,
batch_size=batch_size,
):
Xi, yi = to_device(Xi, device), to_device(yi, device)
y_pred = model(Xi)
y_pred = torch.log(y_pred)
loss = criterion(y_pred, yi)
y_preds.append(y_pred)
loss = loss.item()
losses.append(loss)
batch_sizes.append(len(Xi))
toc = time.time()
return {
'losses': losses,
'batch_sizes': batch_sizes,
'y_proba': torch.cat(y_preds).cpu().detach().numpy(),
'time': toc - tic,
}
def train_step(model, dataset, device, criterion, batch_size, optimizer):
model.train()
y_preds = []
losses = []
batch_sizes = []
tic = time.time()
for Xi, yi in torch.utils.data.DataLoader(dataset, batch_size=batch_size):
Xi, yi = to_device(Xi, device), to_device(yi, device)
optimizer.zero_grad()
y_pred = model(Xi)
y_pred = torch.log(y_pred)
loss = criterion(y_pred, yi)
loss.backward()
optimizer.step()
y_preds.append(y_pred)
losses.append(loss.item())
batch_sizes.append(len(Xi))
toc = time.time()
return {
'losses': losses,
'batch_sizes': batch_sizes,
'y_proba': torch.cat(y_preds).cpu().detach().numpy(),
'time': toc - tic,
}
def test_check_device_torch_tensor(self, to_device, x, device_from,
device_to):
if 'cuda' in (device_from,
device_to) and not torch.cuda.is_available():
pytest.skip()
x = to_device(x, device=device_from)
assert x.device.type == device_from
x = to_device(x, device=device_to)
assert x.device.type == device_to
def test_check_device_torch_tensor(self, to_device, x, device_from, device_to):
if 'cuda' in (device_from, device_to) and not torch.cuda.is_available():
pytest.skip()
prev_device = None
if None in (device_from, device_to):
prev_device = x.device.type
x = to_device(x, device=device_from)
self.check_device_type(x, device_from, prev_device)
x = to_device(x, device=device_to)
self.check_device_type(x, device_to, prev_device)
def test_check_device_packed_padded_sequence(
self, to_device, x_pad_seq, device_from, device_to):
if 'cuda' in (device_from, device_to) and not torch.cuda.is_available():
pytest.skip()
prev_device = None
if None in (device_from, device_to):
prev_device = x_pad_seq.data.device.type
x_pad_seq = to_device(x_pad_seq, device=device_from)
self.check_device_type(x_pad_seq.data, device_from, prev_device)
x_pad_seq = to_device(x_pad_seq, device=device_to)
self.check_device_type(x_pad_seq.data, device_to, prev_device)
def initialize_criterion(self):
"""Initializes the criterion."""
criterion_params = self._get_params_for('criterion')
self.criterion_ = self.criterion(**criterion_params)
if isinstance(self.criterion_, torch.nn.Module):
self.criterion_ = to_device(self.criterion_, self.device)
return self
def test_check_device_tuple_torch_tensor(
self, to_device, x_tup, device_from, device_to):
if 'cuda' in (device_from, device_to) and not torch.cuda.is_available():
pytest.skip()
prev_devices = [None for _ in range(len(x_tup))]
if None in (device_from, device_to):
prev_devices = [x.device.type for x in x_tup]
x_tup = to_device(x_tup, device=device_from)
for xi, prev_d in zip(x_tup, prev_devices):
self.check_device_type(xi, device_from, prev_d)
x_tup = to_device(x_tup, device=device_to)
for xi, prev_d in zip(x_tup, prev_devices):
self.check_device_type(xi, device_to, prev_d)
def batchify(self, data, bsz):
# Work out how cleanly we can divide the dataset into bsz parts.
nbatch = data.size(0) // bsz
# Trim off any extra elements that wouldn't cleanly fit (remainders).
data = data.narrow(0, 0, nbatch * bsz)
# Evenly divide the data across the bsz batches.
data = data.view(bsz, -1).t().contiguous()
return to_device(data, self.device)
def initialize_module(self):
"""Initializes the module.
Note that if the module has learned parameters, those will be
reset.
"""
kwargs = self._get_params_for('module')
kwargs_critic = self._get_params_for('critic')
module = self.module
critic = self.critic
is_initialized = isinstance(module, torch.nn.Module)
is_initialized_critic = isinstance(critic, torch.nn.Module)
if kwargs or not is_initialized:
if is_initialized:
module = type(module)
if (is_initialized or self.initialized_) and self.verbose:
msg = self._format_reinit_msg("module", kwargs)
print(msg)
module = module(**kwargs)
if kwargs_critic or not is_initialized_critic:
if is_initialized_critic:
critic = type(critic)
if (is_initialized_critic or self.initialized_) and self.verbose:
msg = self._format_reinit_msg("critic", kwargs_critic)
print(msg)
critic = self.distance_(critic, **kwargs_critic)
self.module_ = to_device(module, self.device) # type: torch.nn.Module
self.critic_ = to_device(critic, self.device) # type: Critic
return self
def train_torch(
model,
X,
X_test,
y,
y_test,
batch_size,
device,
lr,
max_epochs,
):
model = to_device(model, device)
idx_train, idx_valid = next(
iter(StratifiedKFold(5, random_state=0).split(np.arange(len(X)), y)))
X_train, X_valid, y_train, y_valid = (X[idx_train], X[idx_valid],
y[idx_train], y[idx_valid])
dataset_train = torch.utils.data.TensorDataset(
torch.tensor(X_train),
torch.tensor(y_train),
)
dataset_valid = torch.utils.data.TensorDataset(
torch.tensor(X_valid),
torch.tensor(y_valid),
)
optimizer = torch.optim.Adadelta(model.parameters(), lr=lr)
criterion = nn.NLLLoss()
for epoch in range(max_epochs):
train_out = train_step(
model,
dataset_train,
batch_size=batch_size,
device=device,
criterion=criterion,
optimizer=optimizer,
)
report(y=y_train, epoch=epoch, training=True, **train_out)
valid_out = valid_step(
model,
dataset_valid,
batch_size=batch_size,
device=device,
criterion=criterion,
)
report(y=y_valid, epoch=epoch, training=False, **valid_out)
print('-' * 50)
return model
def test_nested_data(self, to_device, x_list, device_from, device_to):
# Sometimes data is nested because it would need to be padded so it's
# easier to return a list of tensors with different shapes.
# to_device should honor this.
if 'cuda' in (device_from, device_to) and not torch.cuda.is_available():
pytest.skip()
prev_devices = [None for _ in range(len(x_list))]
if None in (device_from, device_to):
prev_devices = [x.device.type for x in x_list]
x_list = to_device(x_list, device=device_from)
assert isinstance(x_list, list)
for xi, prev_d in zip(x_list, prev_devices):
self.check_device_type(xi, device_from, prev_d)
x_list = to_device(x_list, device=device_to)
assert isinstance(x_list, list)
for xi, prev_d in zip(x_list, prev_devices):
self.check_device_type(xi, device_to, prev_d)
def test_check_device_dict_torch_tensor(
self, to_device, x_dict, device_from, device_to):
if 'cuda' in (device_from, device_to) and not torch.cuda.is_available():
pytest.skip()
original_x_dict = deepcopy(x_dict)
prev_devices=[None for _ in range(len(list(x_dict.keys())))]
if None in (device_from, device_to):
prev_devices = [x.device.type for x in x_dict.values()]
new_x_dict = to_device(x_dict, device=device_from)
for xi, prev_d in zip(new_x_dict.values(), prev_devices):
self.check_device_type(xi, device_from, prev_d)
new_x_dict = to_device(new_x_dict, device=device_to)
for xi, prev_d in zip(new_x_dict.values(), prev_devices):
self.check_device_type(xi, device_to, prev_d)
assert x_dict.keys() == original_x_dict.keys()
for k in x_dict:
assert np.allclose(x_dict[k], original_x_dict[k])
def cached_forward_iter(*args, device=net.device, **kwargs):
for yp in y_preds:
yield to_device(yp, device=device)
