Version of densenet

[haofanwang]

May I ask what's the version of densenet in torchivision.models ? The original or efficient_densenet_pytorch, as the original is memory hungry. If it's the original version, would pytorch team consider add the efficient version into model zoo ?

[soumith]

we plan to add a memory efficient version soon via pytorch/pytorch#4594

[gpleiss]

Now that PyTorch 0.4 is officially out, I'm making the efficient_densenet_pytorch code use the checkpointing feature. I can make a PR to this repo once I get it working!

[fmassa]

@gpleiss that would be a nice addition! Maybe by specifying a constructor argument that dispatches to checkpoint? The only thing we need to keep in mind is that checkpoint currently requires the input to have requires_grad=True, which is suboptimal in cases where we don't checkpoint.

[gpleiss]

@fmassa I'm thinking something like this?

# prev_features = [feat_1, feat_2, ...]
# ...
 if self.efficient and any(prev_features.requires_grad):
    bottleneck_output = checkpoint(bn_function, *prev_features)
 else:
    bottleneck_output = bn_function(*prev_features)
# ...

And the self.efficient flag is something that can be passed in by the user.

[fmassa]

That sounds good!

[gpleiss]

I'm profiling it now for my repo, and it seems to be good! Should have a PR ready tomorrow.

[gpleiss]

So sorry for the late response this...

The efficient densenet code seems to work great on a single GPU. However, on multiple GPUs with nn.DataParallel, @wandering007 points out that the checkpointing feature is quite slow (see gpleiss/efficient_densenet_pytorch#36). I think this is because checkpointing requires some sort of inter-GPU synchronization.

I'm opening up an issue in PyTorch about this. I'm holding off on a PR for now.

[fmassa]

Sounds good, thanks @gpleiss !