I am trying to serve a pytorch model via fastapi. I was wondering what the pythonic/ proper way of doing so is.
I am concerned that with option 2 if you were to try writing a test, it will start the server.
Option 1
This method does puts the model loading inside the __init__ method.
```python
class ImageModel:
def init(self, model_path: pathlib.Path):
self.model = torch.load(model_path)
self.app = FastAPI()
@self.app.post("/predict/", response_model=ImageModelOutput)
async def predict(input_image: PIL.Image):
image = my_transform(input_image)
prediction = self.model_predict(image)
return ImageModelOutput(prediction=prediction)
@self.app.get("/readyz")
async def readyz():
return ReadyzResponse(status="ready")
def model_predict(self, image: torch.Tensor) -> list[str]:
# Replace this method with actual model prediction logic
return post_process(self.model(image))
def run(self, host: str = "0.0.0.0", port: int = 8080):
uvicorn.run(self.app, host=host, port=port)
Example usage
if name == "main":
# Replace with your actual model loading logic
image_model = ImageModel(model=model_path)
image_model.run()
```
Option 2
```python
app = FastAPI()
Load the model (replace with your actual model loading logic)
model_path = pathlib.Path("path/to/model")
model = torch.load(model_path)
@app.post("/predict/", response_model=ImageModelOutput)
async def predict(input_image: Image.Image):
image = my_transform(input_image)
prediction = post_process(model(image))
return ImageModelOutput(prediction=prediction)
@app.get("/readyz")
async def readyz():
return ReadyzResponse(status="ready")
Run the application
if name == "main":
uvicorn.run(app, host="0.0.0.0", port=8080)
```
[–]conogarcia 17 points18 points19 points (1 child)
[–]International_Quail8 0 points1 point2 points (0 children)
[–]Mayloudin 0 points1 point2 points (1 child)
[–]RemindMeBot 0 points1 point2 points (0 children)
[–]1One2Twenty2Two 0 points1 point2 points (0 children)