Script for computing memory consumption

[DaShenZi721]

Hello!

In the Figure 1, how do you get memory consumption of cosFormer on the LRA benchmark? Could you please open source script for computing it?

[OpenNLPLab123]

Sorry, the previous script has been lost. Since S4 already provides a torch-based lra codebase, this becomes easier, and here are two ways to get memory consumption. Option 1: Use the following code:

gb_used = torch.cuda.max_memory_allocated() / 1024 / 1024 / 1024
torch.cuda.reset_peak_memory_stats()
gb_free = self.cuda_env.total_memory_in_GB - gb_used
self.total_memory_in_GB = prop.total_memory / 1024 / 1024 / 1024
prop = torch.cuda.get_device_properties("cuda:{}".format(cur_device))

Option 2: Based on https://github.com/Oldpan/Pytorch-Memory-Utils

[DaShenZi721]

Thanks for your reply! If there is a PyTorch-based version, I am also very willing to use it! May I ask where is this repository?

[OpenNLPLab123]

Thanks for your reply! If there is a PyTorch-based version, I am also very willing to use it! May I ask where is this repository?

For the PyTorch-based version, you can use our codebase, which is based on S4, or the origin S4 codebase.