ColossalAI
ColossalAI copied to clipboard
[fx] fix the false interpretation of algorithm 3 in https://arxiv.org/abs/1604.06174.
What's new?
After a discussion with my collaborator @Cypher30, I figured out my misinterpretation of algorithm 3 in https://arxiv.org/abs/1604.06174. This PR can be viewed as an amendment to this previous PR #1433. Now users don't need to give a specific budget, the budget is automatically estimated through grid search.
model = resnet18()
input_sample = torch.rand(4, 3, 224, 224)
gm = symbolic_trace(model)
MetaInfoProp(gm).run(input_sample)
gm = chen_greedy(gm)