search

gordicaleksa / pytorch-neural-style-transfer

Reconstruction of the original paper on neural style transfer (Gatys et al.). I've additionally included reconstruction scripts which allow you to reconstruct only the content or the style of the image - for better understanding of how NST works.
https://youtube.com/c/TheAIEpiphany
MIT License
376 stars 79 forks source link

Support for Nvidia Ampere cards #7

Open entmike opened 3 years ago

entmike commented 3 years ago

Just an FYI, I was playing around with this repo all morning (awesome work!) and was wondering why it was so slow. It turns out that it was using CPU and not CUDA/GPU. (It would be nice if you spammed a warning for others in case they are not as stubborn as me) I am using a newer Ampere card that needs CUDA 11, so I had to solve the issue as shown below. I hope this might help others. Note that this is similar to another closed issue by somebody else, however I had to make further changes such as adding cudatoolkit=11.1 to mine.

name: pytorch-nst
channels:
  - defaults
  - pytorch
  - conda-forge
dependencies:
  - cudatoolkit=11.1
  - python=3.7.6
  - pip=20.0.2
  - matplotlib=3.1.3
  - torchvision
  - torchaudio
  - pip:
    - numpy==1.18.1
    - opencv-python==4.2.0.32
    - conda-forge::pytorch
Finerrkekz commented 2 years ago

Managed to get it to work by running this line in anaconda

pip install torch===1.7.1+cu110 torchvision===0.8.2+cu110 torchaudio===0.7.2 -f https://download.pytorch.org/whl/torch_stable.html