!pip install "thinc>=8.0.0a0" torch "tensorflow>=2.0"

from thinc.api import require_gpu

require_gpu()

"""
Explanation: Using a single memory pool for Cupy and PyTorch or TensorFlow
Requesting memory from a GPU device directly is expensive, so most deep learning libraries will over-allocate, and maintain an internal pool of memory they will keep a hold of, instead of returning it back to the device. This means the libraries don't by default play well together: they all expect to be the single consumer of the GPU memory, so they hog it selfishly. If you use two frameworks together, you can get unexpected out-of-memory errors.
Thinc's internal models use cupy for GPU operations, and cupy offers a nice solution for this problem. You can provide cupy with a custom memory allocation function, which allows us to route cupy's memory requests via another library. This avoids the memory problem when you use PyTorch and cupy together, or when you use cupy and Tensorflow together. We don't yet have a similar solution for using PyTorch and Tensorflow together, however.
To start with, we call the require_gpu() function, which tells Thinc and PyTorch to allocate on GPU.
End of explanation
"""


from thinc.api import use_pytorch_for_gpu_memory

use_pytorch_for_gpu_memory()

"""
Explanation: We then call use_pytorch_for_gpu_memory() to set up the allocation strategy. Now when cupy tries to request GPU memory, it will do so by asking PyTorch, rather than asking the GPU directly.
End of explanation
"""


import cupy 
import torch.cuda

def allocate_cupy_tensor(size):
    array = cupy.zeros((size,), dtype="f")
    print(array.size, torch.cuda.max_memory_allocated())
    return array
allocate_cupy_tensor(16)

"""
Explanation: To test that it's working, we make a little function that allocates an array using cupy, and prints its size, along with the current size of PyTorch's memory pool. Notice the over-allocation: PyTorch grabs a much bigger chunk of memory than just our little array. That's why we need to have only one memory pool.
End of explanation
"""


import tensorflow

with tensorflow.device('/device:GPU:0'):
    arr = allocate_cupy_tensor(1000)
    arr = None
    arr = allocate_cupy_tensor(1000)
    arr = None

"""
Explanation: We can also see that even when we free the tensor, the memory isn't immediately released. On the other hand, we don't need to resize the memory pool when we make a second small allocation.
End of explanation
"""


arr = allocate_cupy_tensor(1000)
for _ in range(100):
    arr2 = allocate_cupy_tensor(900000)

"""
Explanation: If we make a huge allocation, we'll have to resize the pool though. Let's make sure the pool resizes properly, and  that memory is freed when the tensors are removed.
End of explanation
"""


