def test(batch_size=200): model.eval() test_loss = [] idx = np.random.permutation(x_test.shape[0]) count = 0 while count < x_test.shape[0] / batch_size: batch_index = idx[count * batch_size:(count + 1) * batch_size] loss = dsn_loss(model(x_test[batch_index, :]), label_test[batch_index], event_test[batch_index]) test_loss += [loss.data] count += 1 return test_loss[:-1]
def train(batch_size=100): model.train() train_loss = [] idx = np.random.permutation(x_train.shape[0]) count = 0 while count < len(idx) / batch_size: batch_index = idx[count * batch_size:(count + 1) * batch_size] optimizer.zero_grad() loss = dsn_loss(model(x_train[batch_index, :]), label_train[batch_index], event_train[batch_index]) loss.backward() train_loss += [loss.data] optimizer.step() count += 1 return train_loss[:-1]