def test_torch_e2e_state_dict(ray_start_4_cpus):
def train_func():
model = torch.nn.Linear(1, 1).state_dict()
train.save_checkpoint(model=model)
scaling_config = {"num_workers": 2}
trainer = TorchTrainer(train_loop_per_worker=train_func,
scaling_config=scaling_config)
result = trainer.fit()
# If loading from a state dict, a model definition must be passed in.
with pytest.raises(ValueError):
TorchPredictor.from_checkpoint(result.checkpoint)
class TorchScorer:
def __init__(self):
self.pred = TorchPredictor.from_checkpoint(result.checkpoint,
model=torch.nn.Linear(
1, 1))
def __call__(self, x):
return self.pred.predict(x, dtype=torch.float)
predict_dataset = ray.data.range(3)
predictions = predict_dataset.map_batches(TorchScorer,
batch_format="pandas",
compute="actors")
assert predictions.count() == 3
def test_predict_model_not_training(model):
predictor = TorchPredictor(model=model)
data_batch = np.array([1])
predictor.predict(data_batch)
assert not predictor.model.training
def test_predict_feature_columns(model):
predictor = TorchPredictor(model=model)
data_batch = np.array([[1, 4], [2, 5], [3, 6]])
predictions = predictor.predict(data_batch, feature_columns=[0])
assert len(predictions) == 3
assert predictions.to_numpy().flatten().tolist() == [2, 4, 6]
def test_predict_with_preprocessor(model, preprocessor):
predictor = TorchPredictor(model=model, preprocessor=preprocessor)
data_batch = np.array([[1], [2], [3]])
predictions = predictor.predict(data_batch)
assert len(predictions) == 3
assert predictions.to_numpy().flatten().tolist() == [4, 8, 12]
def test_predict_dataframe():
predictor = TorchPredictor(model=torch.nn.Linear(2, 1, bias=False))
data_batch = pd.DataFrame({"X0": [0.0, 0.0, 0.0], "X1": [0.0, 0.0, 0.0]})
predictions = predictor.predict(data_batch, dtype=torch.float)
assert len(predictions) == 3
assert predictions.to_numpy().flatten().tolist() == [0.0, 0.0, 0.0]
def test_predict_array(model):
predictor = TorchPredictor(model=model)
data_batch = np.array([[1], [2], [3]])
predictions = predictor.predict(data_batch)
assert len(predictions) == 3
assert predictions.to_numpy().flatten().tolist() == [2, 4, 6]
def test_predict_dataframe_with_feature_columns():
predictor = TorchPredictor(model=torch.nn.Identity())
data_batch = pd.DataFrame({"X0": [0.0, 0.0, 0.0], "X1": [1.0, 1.0, 1.0]})
predictions = predictor.predict(data_batch, feature_columns=["X0"])
assert len(predictions) == 3
assert predictions.to_numpy().flatten().tolist() == [0.0, 0.0, 0.0]
def test_predict_array_with_different_dtypes(input_dtype,
expected_output_dtype):
predictor = TorchPredictor(model=torch.nn.Identity())
data_batch = np.array([[1], [2], [3]])
predictions = predictor.predict(data_batch, dtype=input_dtype)
assert all(prediction.dtype == expected_output_dtype
for prediction in predictions["predictions"])
def test_init(model, preprocessor):
predictor = TorchPredictor(model=model, preprocessor=preprocessor)
checkpoint = {MODEL_KEY: model, PREPROCESSOR_KEY: preprocessor}
checkpoint_predictor = TorchPredictor.from_checkpoint(
Checkpoint.from_dict(checkpoint))
assert checkpoint_predictor.model == predictor.model
assert checkpoint_predictor.preprocessor == predictor.preprocessor
def test_predict_array_output(model):
"""Tests if predictor works if model outputs an array instead of single value."""
predictor = TorchPredictor(model=model)
data_batch = np.array([[1, 1], [2, 2], [3, 3]])
predictions = predictor.predict(data_batch)
assert len(predictions) == 3
assert np.array_equal(
predictions.to_numpy().flatten().tolist(), [[2, 2], [4, 4], [6, 6]]
)
def test_horovod_state_dict(ray_start_4_cpus):
def train_func(config):
result = hvd_train_func(config)
assert len(result) == epochs
assert result[-1] < result[0]
num_workers = 2
epochs = 10
scaling_config = {"num_workers": num_workers}
config = {"num_epochs": epochs, "save_model_as_dict": True}
trainer = HorovodTrainer(
train_loop_per_worker=train_func,
train_loop_config=config,
scaling_config=scaling_config,
)
result = trainer.fit()
predictor = TorchPredictor.from_checkpoint(result.checkpoint, model=Net())
# Find some test data to run on.
test_set = datasets.MNIST(
"./data",
train=False,
download=True,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307, ), (0.3081, ))
]),
)
test_dataloader = DataLoader(test_set, batch_size=10)
test_dataloader_iter = iter(test_dataloader)
images, labels = next(
test_dataloader_iter) # only running a batch inference of 10 images
predicted_labels = run_image_prediction(predictor.model, images)
assert torch.equal(predicted_labels, labels)
def test_predict_from_checkpoint_no_preprocessor(model):
checkpoint = Checkpoint.from_dict({MODEL_KEY: model})
predictor = TorchPredictor.from_checkpoint(checkpoint)
data_batch = np.array([[1], [2], [3]])
predictions = predictor.predict(data_batch)
assert len(predictions) == 3
assert predictions.to_numpy().flatten().tolist() == [2, 4, 6]
def __init__(self):
self.pred = TorchPredictor.from_checkpoint(result.checkpoint,
model=torch.nn.Linear(
1, 1))
def __init__(self):
self.pred = TorchPredictor.from_checkpoint(result.checkpoint)
def __init__(self):
checkpoint = Checkpoint.from_object_ref(checkpoint_object_ref)
self.predictor = TorchPredictor.from_checkpoint(checkpoint)