Python cog 관련

Python cog 관련

 

 

 

 

 

https://pypi.org/project/cog/

cog

 

cog는 머신 러닝을 위한 컨테이너를 지칭하는 패키지 이름

 

 

우선 제공 모델을 돌리기 위한 Docker 환경을 cog.yaml에 정의한다.

 

predict: predict.py:Predictor
build:
  gpu: true
  python_version: 3.8
  python_packages:
    - torch==1.8.1
    - torchvision==0.8.1
    - lmdb==1.1.1
    - pillow==8.1.2

 

predict.py에 어떻게 prediction에서 모델을 돌릴지 정의

 

from cog import BasePredictor, Input, Path
import torch

class Predictor(BasePredictor):
    def setup(self):
        """Load the model into memory to make running multiple predictions efficient"""
        self.model = torch.load("./weights.pth")

    # The arguments and types the model takes as input
    def predict(self,
          image: Path = Input(title="Grayscale input image")
    ) -> Path:
        """Run a single prediction on the model"""
        processed_image = preprocess(image)
        output = self.model(processed_image)
        return postprocess(output)

 

참조

https://github.com/rosinality/style-based-gan-pytorch/blob/master/predict.py

GitHub - rosinality/style-based-gan-pytorch: Implementation A Style-Based Generator Architecture for Generative Adversarial Netw