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_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_multi_modal_real_model(use_gpu):
class CustomModule(torch.nn.Module):
def __init__(self):
super().__init__()
self.linear1 = torch.nn.Linear(1, 1)
self.linear2 = torch.nn.Linear(1, 1)
def forward(self, input_dict: dict):
out1 = self.linear1(input_dict["A"])
out2 = self.linear2(input_dict["B"])
return out1 + out2
predictor = TorchPredictor(model=CustomModule(), use_gpu=use_gpu)
data = pd.DataFrame([[1, 2], [3, 4]], columns=["A", "B"])
predictions = predictor.predict(data, dtype=torch.float)
assert len(predictions) == 2
if use_gpu:
assert next(
predictor.model.parameters()
).is_cuda, "Model should be moved to GPU if use_gpu is True"
else:
assert not next(
predictor.model.parameters()
).is_cuda, "Model should not be on GPU if use_gpu is False"
def test_array_real_model():
model = torch.nn.Linear(2, 1)
predictor = TorchPredictor(model=model)
data = np.array([[1, 2], [3, 4]])
predictions = predictor.predict(data, dtype=torch.float)
assert len(predictions) == 2
def test_predict_array(model):
predictor = TorchPredictor(model=model)
data_batch = np.asarray([[1], [2], [3]])
predictions = predictor.predict(data_batch)
assert len(predictions) == 3
assert predictions.flatten().tolist() == [2, 4, 6]
def test_predict_array_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.flatten().tolist() == [4, 8, 12]
def test_predict_array_with_preprocessor(model, preprocessor, use_gpu):
predictor = TorchPredictor(model=model, preprocessor=preprocessor, use_gpu=use_gpu)
data_batch = np.array([1, 2, 3])
predictions = predictor.predict(data_batch)
assert len(predictions) == 3
assert predictions.flatten().tolist() == [4, 8, 12]
def test_predict_array(model, use_gpu):
predictor = TorchPredictor(model=model, use_gpu=use_gpu)
data_batch = np.asarray([1, 2, 3])
predictions = predictor.predict(data_batch)
assert len(predictions) == 3
assert predictions.flatten().tolist() == [2, 4, 6]
def test_predict_dataframe():
predictor = TorchPredictor(model=DummyModelMultiInput())
data_batch = pd.DataFrame({"X0": [0.0, 0.0, 0.0], "X1": [1.0, 2.0, 3.0]})
predictions = predictor.predict(data_batch, dtype=torch.float)
assert len(predictions) == 3
assert predictions.to_numpy().flatten().tolist() == [1.0, 2.0, 3.0]
def test_predict_array_with_different_dtypes(model, input_dtype,
expected_output_dtype):
predictor = TorchPredictor(model=model)
data_batch = np.array([[1], [2], [3]])
predictions = predictor.predict(data_batch, dtype=input_dtype)
assert predictions.dtype == expected_output_dtype
def test_predict_array_with_different_dtypes(
model, input_dtype, expected_output_dtype, use_gpu
):
predictor = TorchPredictor(model=model, use_gpu=use_gpu)
data_batch = np.array([1, 2, 3])
predictions = predictor.predict(data_batch, dtype=input_dtype)
assert predictions.dtype == expected_output_dtype
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(batch_type):
predictor = TorchPredictor(model=DummyModelMultiInput())
raw_batch = pd.DataFrame({"X0": [0.0, 0.0, 0.0], "X1": [1.0, 2.0, 3.0]})
data_batch = convert_pandas_to_batch_type(raw_batch, type=TYPE_TO_ENUM[batch_type])
raw_predictions = predictor.predict(data_batch, dtype=torch.float)
predictions = convert_batch_type_to_pandas(raw_predictions)
assert len(predictions) == 3
assert predictions.to_numpy().flatten().tolist() == [1.0, 2.0, 3.0]
def test_predict_multi_output():
predictor = TorchPredictor(model=DummyModelMultiOutput())
data_batch = np.array([[1], [2], [3]])
predictions = predictor.predict(data_batch)
# Model outputs two tensors
assert len(predictions) == 2
for k, v in predictions.items():
# Each tensor is of size 3
assert len(v) == 3
assert v.flatten().tolist() == [1, 2, 3]
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 = ScalingConfig(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_multi_modal_real_model():
class CustomModule(torch.nn.Module):
def __init__(self):
super().__init__()
self.linear1 = torch.nn.Linear(1, 1)
self.linear2 = torch.nn.Linear(1, 1)
def forward(self, input_dict: dict):
out1 = self.linear1(input_dict["A"])
out2 = self.linear2(input_dict["B"])
return out1 + out2
predictor = TorchPredictor(model=CustomModule())
data = pd.DataFrame([[1, 2], [3, 4]], columns=["A", "B"])
predictions = predictor.predict(data, dtype=torch.float)
assert len(predictions) == 2
def test_predict_array_no_training(model, use_gpu):
checkpoint = TorchCheckpoint.from_model(model)
predictor = TorchPredictor.from_checkpoint(checkpoint, use_gpu=use_gpu)
data_batch = np.array([1, 2, 3])
predictions = predictor.predict(data_batch)
assert len(predictions) == 3
assert predictions.flatten().tolist() == [2, 4, 6]
def test_predict_array_no_training(model):
checkpoint = to_air_checkpoint(model)
predictor = TorchPredictor.from_checkpoint(checkpoint)
data_batch = np.array([[1], [2], [3]])
predictions = predictor.predict(data_batch)
assert len(predictions) == 3
assert predictions.flatten().tolist() == [2, 4, 6]
def __init__(self):
self.pred = TorchPredictor.from_checkpoint(result.checkpoint)
def __init__(self):
self.pred = TorchPredictor.from_checkpoint(result.checkpoint,
model=torch.nn.Linear(
1, 1))
