a way to reduce vram usage 2.7 times, from @lmxyy at mit-han-lab

[frankyifei]

Should be easy to implement, reduce vram usage up to 2.7 times,

All you need to do is use convert_model to wrap all the Conv2d in your PyTorch model to our PatchConv. For example,

from patch_conv import convert_model

model = Model(...)  # Your PyTorch model
model = convert_model(model, splits=4)  # The only modification you need to make

with torch.no_grad():
    model(...)  # Run the model in the original way

link