search

lucidrains / enformer-pytorch

Implementation of Enformer, Deepmind's attention network for predicting gene expression, in Pytorch
MIT License
435 stars 82 forks source link

Make the output activation in HeadAdapterWrapper adjustable #23

Closed gokceneraslan closed 1 year ago

gokceneraslan commented 1 year ago

So that different heads can be used for different type of predictions and/or loss functions

lucidrains commented 1 year ago

@gokceneraslan hey! thanks for the PR! do you want to see if 0.7.1 works for your use-case? (i also updated the other fine tuning wrappers with similar logic)

are you seeing better success finetuning with a different activation function, and which one?

gokceneraslan commented 1 year ago

Wow thanks for lightning fast response and version bump! 💯

I am just playing around with bulk ATAC-seq models with binarized predictions i.e. sigmoid+bce, see below for a nice abuse of this feature:

model = Enformer.from_pretrained('EleutherAI/enformer-official-rough', target_length=-1)
head_out = nn.Sequential(
    Rearrange('b t c -> b c t'),
    nn.AvgPool1d(8),
    Rearrange('... () -> ...'),
    nn.Sigmoid()
)
model = HeadAdapterWrapper(enformer=model, num_tracks=labels.shape[1], output_activation=head_out)
model = model.to('cuda:1')
lucidrains commented 1 year ago

Oh noice :^)

wejarrard commented 1 year ago

@gokceneraslan Did this end up working well for you to do binarized predictions? were there any other hoops you needed to jump through?