[D] Tensorflow vs. PyTorch Memory Usage

Submitted by Oceanboi t3_zm6h07 in MachineLearning

Why is it that when I go to create a CNN with 4 layers (output channels: 64, 32, 16, 16), I can do this in PyTorch, but in Tensorflow I get resource errors saying I don't have enough resources?

For reference I am using a stock NVIDIA RTX 3080.

Also, now that I am experimenting with larger models, would I benefit from renting TPU? Does this make the actual models train faster and would it help with larger batches?

3

Comments

You must log in or register to comment.

mofawzy89 t1_j09kc7w wrote

I'm not sure about memory management between both but I faced thr same for BiLSTM For large models use gcp better yes tpus or nvidia A100

1

schludy t1_j09ukmr wrote

Do you handle the data the same way? Maybe you're loading more data in the tensorflow implementation. It's really hard to tell without seeing the code

6

MrFlufypants t1_j0ao1ei wrote

I’ve had issues where tensorflow automatically grabs the whole gpu while PyTorch only uses what the model asks for. Could totally not be your problem, but if you’re running multiple models it could be your problem

1

MrFlufypants t1_j0atiwu wrote

Replying to veb101 (#934,930)

There are a couple ways to do it. That’s the one I use normally. Sometimes that doesn’t work though. Can’t quite remember the use case where it wasn’t working

1

Oceanboi OP t1_j0cpolp wrote

Replying to schludy (#933,481)

Data is handled both the same way. I think it has to do with what u/MrFlufypants said, because when I restart my kernel and run it after freeing up some resources, it runs. I think the number of filters I was setting are right at the threshold in which my GPU runs out of VRAM, so small memory management differences in TF and PyTorch are causing TF to hit the limit faster than PyTorch.

1