Upload 31 files

cop/hic_model.py

@@ -0,0 +1,207 @@
+import os,sys
+from pretrain.track.layers import AttentionPool,CNN
+from pretrain.track.transformers import Transformer
+from einops.layers.torch import Rearrange
+from einops import rearrange
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+import torchvision.transforms as T
+
+class Convblock(nn.Module):
+    def __init__(self,in_channel,kernel_size,dilate_size,dropout=0.1):
+        super().__init__()
+        self.conv=nn.Sequential(
+            nn.Conv2d(
+                in_channel, in_channel,
+                kernel_size, padding=self.pad(kernel_size, dilate_size),
+                dilation=dilate_size),
+            nn.GroupNorm(16, in_channel),
+            nn.Dropout(dropout)
+        )
+    def pad(self,kernelsize, dialte_size):
+        return (kernelsize - 1) * dialte_size // 2
+    def symmetric(self,x):
+        return (x + x.permute(0,1,3,2)) / 2
+    def forward(self,x):
+        identity=x
+        out=self.conv(x)
+        x=out+identity
+        x=self.symmetric(x)
+        return F.relu(x)
+
+class dilated_tower(nn.Module):
+    def __init__(self,embed_dim,in_channel=48,kernel_size=9,dilate_rate=4):
+        super().__init__()
+        dilate_convs=[]
+        for i in range(dilate_rate+1):
+            dilate_convs.append(
+                Convblock(in_channel,kernel_size=kernel_size,dilate_size=2**i))
+
+        self.cnn=nn.Sequential(
+            Rearrange('b l n d -> b d l n'),
+            nn.Conv2d(embed_dim, in_channel, kernel_size=1),
+            *dilate_convs,
+            nn.Conv2d(in_channel, in_channel, kernel_size=1),
+            Rearrange('b d l n -> b l n d'),
+        )
+    def forward(self,x,crop):
+        x=self.cnn(x)
+        x=x[:,crop:-crop,crop:-crop,:]
+        return x
+
+class discriminator(nn.Module):
+    def __init__(self,):
+        super(discriminator, self).__init__()
+        self.conv=nn.Sequential(
+            nn.Conv2d(1,32,kernel_size=7,padding=3),
+            nn.ReLU(),
+            nn.MaxPool1d(4),
+            nn.Conv2d(32,64,kernel_size=7,padding=3),
+            nn.ReLU(),
+            nn.MaxPool1d(4),
+            nn.Conv2d(64, 64, kernel_size=1),
+        )
+        self.linear=nn.Linear(64,1)
+
+    def complete_mat(self,x,smooth=True):
+        tmp=x.copy()
+        tmp.fill_diagonal_(0)
+        if smooth:
+            t = T.GaussianBlur(kernel_size=5, sigma=0.5)
+            tmp = t(tmp.T + x)
+        else:
+            tmp=tmp.T + x
+        return tmp
+    def forward(self,x1,x2):
+        x1,x2=self.complete_mat(x1),self.complete_mat(x2)
+        diff_mat=(x1-x2)**2
+        diff_mat=self.conv(diff_mat).mean(dim=(-2, -1))
+        return self.linear(diff_mat)
+
+
+
+class Tranmodel(nn.Module):
+    def __init__(self, backbone, transfomer):
+        super().__init__()
+        self.backbone = backbone
+        self.transformer = transfomer
+        hidden_dim = transfomer.d_model
+        self.input_proj = nn.Conv1d(backbone.num_channels, hidden_dim, kernel_size=1)
+    def forward(self, input):
+        input=rearrange(input,'b n c l -> (b n) c l')
+        src = self.backbone(input)
+        src=self.input_proj(src)
+        src = self.transformer(src)
+        return src
+
+class finetunemodel(nn.Module):
+    def __init__(
+        self,
+        pretrain_model,
+        hidden_dim,
+        embed_dim,
+        device,
+        bins=200,
+        in_dim=64,
+        max_bin=10,
+        crop=4,
+        output_dim=1
+    ):
+        super().__init__()
+        self.pretrain_model = pretrain_model
+        self.bins = bins
+        self.max_bin = max_bin
+        self.attention_pool = AttentionPool(hidden_dim)
+        self.crop = crop
+        self.project = nn.Sequential(
+            Rearrange('(b n) c -> b c n', n=bins * 5),
+            nn.Conv1d(hidden_dim, hidden_dim, kernel_size=15, padding=7, groups=hidden_dim),
+            nn.InstanceNorm1d(hidden_dim, affine=True),
+            nn.Conv1d(hidden_dim, embed_dim, kernel_size=1),
+            nn.ReLU(inplace=True),
+            nn.Dropout(0.2)
+        )
+
+        self.cnn = nn.Sequential(
+            nn.Conv1d(embed_dim, embed_dim, kernel_size=15, padding=7),
+            nn.GroupNorm(32, embed_dim),
+            nn.MaxPool1d(kernel_size=5, stride=5),
+            nn.ReLU(inplace=True),
+            nn.Conv1d(embed_dim, embed_dim, kernel_size=1),
+            nn.Dropout(0.2),
+            Rearrange('b c n -> b n c')
+        )
+        encoder_layer = nn.TransformerEncoderLayer(d_model=embed_dim, nhead=4, dim_feedforward=2 * embed_dim,
+                                                   batch_first=True, norm_first=True)
+        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=4)
+
+        self.distance_embed = nn.Embedding(max_bin + 1, embed_dim)
+
+        self.dilate_tower = dilated_tower(embed_dim=embed_dim, in_channel=in_dim)
+        self.prediction_head1 = nn.Linear(in_dim, output_dim)
+        self.dist_dropout = nn.Dropout(0.1)
+        self.device = device
+
+    def position_matrix(self, bins, b, maxbin):
+        pos1 = np.tile(np.arange(bins), (bins, 1))
+        pos2 = pos1.T
+        pos = np.abs(pos1 - pos2)
+        pos = np.where(pos > maxbin, maxbin, pos)
+        pos = np.tile(pos, (b, 1, 1))
+        return torch.tensor(pos).long().to(self.device)
+
+    def upper_tri(self, x, bins):
+        triu_tup = np.triu_indices(bins)
+        d = np.array(list(triu_tup[1] + bins * triu_tup[0]))
+        return x[:, d, :]
+
+    def output_head(self, x, dist_embed, bins):
+        x1 = torch.tile(x.unsqueeze(1), (1, bins, 1, 1))
+        x2 = x1.permute(0, 2, 1, 3)
+        mean_out = (x1 + x2) / 2
+        dot_out = x1 * x2
+        return mean_out + dot_out + dist_embed
+
+    def forward(self, x):
+        b = x.shape[0]
+        x = self.pretrain_model(x)
+        x = self.attention_pool(x)
+        x = self.project(x)
+        x = self.cnn(x)
+        x = self.transformer(x)
+        dist_embed = self.dist_dropout(self.distance_embed(self.position_matrix(self.bins, b=b, maxbin=self.max_bin)))
+        x = self.output_head(x, dist_embed, self.bins)
+        x = self.dilate_tower(x, self.crop)
+        x = rearrange(x, 'b l n d -> b (l n) d')
+        x = self.upper_tri(x, self.bins - 2 * self.crop)
+        x = self.prediction_head1(x)
+        return x
+
+def build_backbone():
+    model = CNN()
+    return model
+def build_transformer(args):
+    return Transformer(
+        d_model=args.hidden_dim,
+        dropout=args.dropout,
+        nhead=args.nheads,
+        dim_feedforward=args.dim_feedforward,
+        num_encoder_layers=args.enc_layers,
+        num_decoder_layers=args.dec_layers
+    )
+
+
+def build_hic_model(args):
+    backbone = build_backbone()
+    transformer = build_transformer(args)
+    pretrain_model = Tranmodel(
+            backbone=backbone,
+            transfomer=transformer,
+        )
+
+    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+    model=finetunemodel(pretrain_model,hidden_dim=args.hidden_dim,embed_dim=args.embed_dim,device=device,bins=args.bins,crop=args.crop,output_dim=3)
+    return model

cop/micro_model.py

@@ -0,0 +1,111 @@
+import os,sys
+import math
+from pretrain.track.model import build_track_model
+import torch.nn as nn
+import torch
+import torch.nn.functional as F
+from einops.layers.torch import Rearrange
+from einops import rearrange
+import numpy as np
+
+class Convblock(nn.Module):
+    def __init__(self,in_channel,kernel_size,dilate_size,dropout=0.1):
+        super().__init__()
+        self.conv=nn.Sequential(
+            nn.Conv2d(
+                in_channel, in_channel,
+                kernel_size, padding=self.pad(kernel_size,1)),
+            nn.GroupNorm(16, in_channel),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Conv2d(
+                in_channel, in_channel,
+                kernel_size, padding=self.pad(kernel_size, dilate_size),
+                dilation=dilate_size),
+        )
+    def pad(self,kernelsize, dialte_size):
+        return (kernelsize - 1) * dialte_size // 2
+    def symmetric(self,x):
+        return (x + x.permute(0,1,3,2)) / 2
+    def forward(self,x):
+        identity=x
+        out=self.conv(x)
+        x=out+identity
+        x=self.symmetric(x)
+        return F.relu(x)
+
+class dilated_tower(nn.Module):
+    def __init__(self,embed_dim,in_channel=64,kernel_size=7,dilate_rate=5):
+        super().__init__()
+        dilate_convs=[]
+        for i in range(dilate_rate+1):
+            dilate_convs.append(
+                Convblock(in_channel,kernel_size=kernel_size,dilate_size=2**i))
+
+        self.cnn=nn.Sequential(
+            Rearrange('b l n d -> b d l n'),
+            nn.Conv2d(embed_dim, in_channel, kernel_size=1),
+            *dilate_convs,
+            nn.Conv2d(in_channel, in_channel, kernel_size=1),
+            Rearrange('b d l n -> b l n d'),
+        )
+    def forward(self,x,crop):
+        x=self.cnn(x)
+        x=x[:,crop:-crop,crop:-crop,:]
+        return x
+
+class Downstream_microc_model(nn.Module):
+    def __init__(
+        self,
+        pretrain_model,
+        embed_dim,
+        hidden_dim=256,
+        in_dim=64,
+        crop=10,
+    ):
+        super().__init__()
+        self.project = nn.Sequential(
+            nn.Linear(embed_dim, 512),
+            nn.ReLU(),
+            nn.Linear(512, hidden_dim),
+        )
+
+        self.pretrain_model=pretrain_model
+        self.dilate_tower = dilated_tower(embed_dim=hidden_dim, in_channel=in_dim,dilate_rate=5)
+        self.prediction_head = nn.Linear(in_dim, 1)
+        self.crop=crop
+
+    def output_head(self, x):
+        bins=x.shape[1]
+        x1 = torch.tile(x.unsqueeze(1), (1, bins, 1, 1))
+        x2 = x1.permute(0, 2, 1, 3)
+        mean_out = (x1 + x2) / 2
+        dot_out = (x1 * x2)/math.sqrt(x.shape[-1])
+        return mean_out + dot_out
+    def upper_tri(self, x,bins):
+        triu_tup = np.triu_indices(bins)
+        d = np.array(list(triu_tup[1] + bins * triu_tup[0]))
+        return x[:, d, :]
+
+    def forward(self,x):
+        x=self.pretrain_model(x)
+        x=self.project(x)
+        x = self.output_head(x)
+        x = self.dilate_tower(x, self.crop)
+        bins = x.shape[1]
+        x = rearrange(x, 'b l n d -> b (l n) d')
+        x = self.upper_tri(x,bins)
+        x = self.prediction_head(x)
+        return x
+
+
+
+def build_microc_model(args):
+    pretrain_model=build_track_model(args)
+    model=Downstream_microc_model(
+        pretrain_model=pretrain_model,
+        embed_dim=args.embed_dim,
+        crop=args.crop
+    )
+    return model
+

data/black_list.bed

@@ -0,0 +1,910 @@
+chr1	628903	635104
+chr1	5850087	5850571
+chr1	8909610	8910014
+chr1	9574580	9574997
+chr1	32043823	32044203
+chr1	33818964	33819344
+chr1	38674335	38674715
+chr1	50017081	50017546
+chr1	52996949	52997329
+chr1	55372488	55372869
+chr1	67971776	67972156
+chr1	73258720	73259100
+chr1	76971068	76971595
+chr1	93936365	93936747
+chr1	93937447	93937827
+chr1	102160407	102160787
+chr1	103620975	103621378
+chr1	106803432	106803816
+chr1	106804021	106804224
+chr1	106804753	106805343
+chr1	121609948	125063427
+chr1	125166231	125184683
+chr1	143184599	143276861
+chr1	146992422	146992802
+chr1	158449073	158449453
+chr1	158872114	158872494
+chr1	159295111	159295493
+chr1	169473895	169474338
+chr1	170006204	170006584
+chr1	172710350	172710732
+chr1	181422611	181423158
+chr1	191961694	191962163
+chr1	195288048	195288429
+chr1	199487949	199488149
+chr1	214709795	214710175
+chr1	215499615	215500014
+chr1	226652017	226652398
+chr1	227699752	227700133
+chr1	229019365	229019745
+chr1	233139985	233140365
+chr1	235520204	235520404
+chr1	235537405	235537785
+chr1	235538899	235540112
+chr1	235540243	235540623
+chr1	235540886	235541649
+chr1	235870625	235871005
+chr1	237940595	237940979
+chr1	237941045	237941514
+chr1	237941893	237942746
+chr1	237943028	237943416
+chr1	237943490	237945232
+chr1	237945285	237946507
+chr1	237948983	237949365
+chr1	237951294	237951802
+chr10	2235555	2235756
+chr10	19746628	19747247
+chr10	19747314	19748342
+chr10	25638376	25638756
+chr10	26873147	26873538
+chr10	30565118	30565501
+chr10	36432964	36433344
+chr10	36434047	36435188
+chr10	37600616	37601002
+chr10	37601246	37601787
+chr10	37601884	37602850
+chr10	38481300	38596500
+chr10	38782600	38967900
+chr10	39000365	41916630
+chr10	42066792	42104971
+chr10	45577925	45578305
+chr10	46706229	46706611
+chr10	47633790	47634172
+chr10	55597861	55600059
+chr10	55626794	55627174
+chr10	57668682	57669062
+chr10	59261793	59262173
+chr10	69590538	69590738
+chr10	69591475	69591858
+chr10	69592355	69592740
+chr10	69592776	69593482
+chr10	69594378	69594760
+chr10	69595141	69595573
+chr10	69595681	69596061
+chr10	77166388	77166768
+chr10	79411056	79411468
+chr10	89786504	89786889
+chr10	100057235	100058064
+chr10	112894488	112894870
+chr10	115056512	115056712
+chr10	123032371	123032751
+chr10	125819621	125820001
+chr10	133689373	133689523
+chr11	8023287	8023667
+chr11	10507706	10510499
+chr11	10768339	10768719
+chr11	10815184	10815384
+chr11	24839563	24839944
+chr11	27850562	27850942
+chr11	47323881	47324333
+chr11	50424039	50813393
+chr11	51081363	54424064
+chr11	64187168	64187556
+chr11	65069483	65069863
+chr11	73510500	73510992
+chr11	81551734	81551934
+chr11	81553835	81554282
+chr11	81556152	81556537
+chr11	81556717	81557101
+chr11	87813427	87814320
+chr11	87815683	87816063
+chr11	103270627	103271007
+chr11	103403270	103403650
+chr11	103404014	103404527
+chr11	103404779	103405289
+chr11	103405809	103406376
+chr11	103406653	103407036
+chr11	103407110	103407310
+chr11	103408089	103409893
+chr11	103410074	103411211
+chr11	110876919	110877308
+chr11	114021166	114021546
+chr11	123003425	123003857
+chr11	123139919	123140301
+chr12	9923	10481
+chr12	2539174	2539982
+chr12	3887955	3888335
+chr12	19795477	19795864
+chr12	20769413	20769432
+chr12	21052950	21053330
+chr12	22005655	22006093
+chr12	27925108	27925488
+chr12	31247541	31247923
+chr12	31247963	31248343
+chr12	34665177	37429869
+chr12	40286245	40286625
+chr12	41363462	41363903
+chr12	41698591	41698971
+chr12	41699048	41699573
+chr12	49817252	49817634
+chr12	62773865	62774257
+chr12	80623983	80624183
+chr12	101486970	101487350
+chr12	123053921	123054301
+chr12	126583199	126583772
+chr12	126584137	126584530
+chr12	130315425	130315904
+chr13	16226300	18171400
+chr13	25984718	25984918
+chr13	31866923	31867303
+chr13	33516898	33517278
+chr13	36065385	36065836
+chr13	40768206	40768595
+chr13	53891451	53891831
+chr13	55971453	55971922
+chr13	56688341	56688749
+chr13	72344211	72344591
+chr13	75592084	75592468
+chr13	83688313	83688693
+chr13	84521524	84522274
+chr13	84522848	84523233
+chr13	88308157	88308357
+chr13	95692549	95692935
+chr13	95693013	95693215
+chr13	95694449	95695698
+chr13	95696145	95696512
+chr13	105488067	105488448
+chr13	107058662	107059042
+chr13	109423944	109424560
+chr14	16000600	18173660
+chr14	23426306	23426691
+chr14	32483953	32485298
+chr14	37490106	37490486
+chr14	40643840	40644220
+chr14	43116742	43117122
+chr14	45238635	45239016
+chr14	45430378	45430758
+chr14	46048457	46048837
+chr14	46847040	46847420
+chr14	51587295	51587847
+chr14	83587331	83587894
+chr14	83588229	83589060
+chr14	84171262	84171729
+chr14	84171838	84172846
+chr14	84173508	84173969
+chr14	84174279	84174691
+chr14	86498937	86499317
+chr15	17058500	19838644
+chr15	30477565	30477945
+chr15	32529779	32530159
+chr15	34715310	34715692
+chr15	35396110	35396495
+chr15	40133887	40134759
+chr15	41157028	41157408
+chr15	52100391	52100771
+chr15	54583731	54584111
+chr15	58152409	58153114
+chr15	58153292	58153690
+chr15	58155859	58156155
+chr15	67040730	67041122
+chr15	91960163	91960543
+chr16	3367430	3368546
+chr16	3369658	3370039
+chr16	3370150	3370542
+chr16	3370932	3371445
+chr16	3371688	3372222
+chr16	10719290	10720105
+chr16	10720417	10720781
+chr16	10721235	10721874
+chr16	10721900	10722280
+chr16	10723423	10723623
+chr16	10723815	10724200
+chr16	10724415	10724654
+chr16	20720929	20721312
+chr16	20721365	20721746
+chr16	20722103	20722552
+chr16	34071571	34071629
+chr16	34131996	34289269
+chr16	34571482	34597852
+chr16	34661168	34661267
+chr16	34919141	34919184
+chr16	35966577	38269112
+chr16	38275767	38280684
+chr16	46380676	46381095
+chr16	46386376	46386491
+chr16	46388622	46389053
+chr16	46390180	46390788
+chr16	46394471	46395088
+chr16	46398828	46401647
+chr16	60470624	60471006
+chr16	65701465	65701846
+chr16	67590312	67590692
+chr16	69358523	69358990
+chr16	73161120	73161500
+chr16	82119745	82120125
+chr17	141682	142062
+chr17	14171308	14171688
+chr17	15568187	15568567
+chr17	19597515	19597985
+chr17	19598613	19599532
+chr17	19599799	19600210
+chr17	19600300	19602064
+chr17	19602160	19602545
+chr17	19602886	19603595
+chr17	19603847	19604047
+chr17	19604922	19605588
+chr17	20851029	20851409
+chr17	21851150	21992060
+chr17	22519042	22520149
+chr17	22520322	22521025
+chr17	22521116	22526407
+chr17	22526636	22530152
+chr17	22530381	22532156
+chr17	22532315	22532940
+chr17	22551066	22551446
+chr17	22813591	26716670
+chr17	26885752	26885795
+chr17	35654769	35655182
+chr17	43251640	43251763
+chr17	43309853	43310048
+chr17	43315021	43316491
+chr17	43997535	43997957
+chr17	53105552	53106565
+chr17	54902920	54903301
+chr17	59279406	59279787
+chr17	63076394	63076777
+chr17	63393238	63393438
+chr17	65555244	65555624
+chr17	72316258	72316638
+chr17	80617407	80617802
+chr18	2842087	2842534
+chr18	8103913	8104113
+chr18	8846332	8846713
+chr18	15457976	20865732
+chr18	34571460	34571840
+chr18	47853089	47853617
+chr18	52883627	52884007
+chr18	59288306	59288686
+chr18	61874562	61874960
+chr18	77455900	77456280
+chr19	246899	247452
+chr19	12105016	12105399
+chr19	13362989	13363369
+chr19	24182199	27257542
+chr19	27741787	27741868
+chr19	36271917	36272148
+chr19	37572465	37572846
+chr19	37576134	37576516
+chr19	46122944	46123324
+chr19	47941356	47941426
+chr19	54794749	54795129
+chr19	56691535	56691736
+chr19	56922158	56922601
+chr2	638427	638808
+chr2	1087103	1087484
+chr2	16271753	16272134
+chr2	22316878	22317258
+chr2	24644617	24644997
+chr2	32916201	32916632
+chr2	33767290	33767703
+chr2	33964664	33965045
+chr2	36276769	36277149
+chr2	40784787	40785278
+chr2	49229452	49230058
+chr2	50588765	50589566
+chr2	54451654	54452034
+chr2	57648677	57649057
+chr2	67953669	67954049
+chr2	75063567	75063994
+chr2	81666317	81666849
+chr2	82814941	82815321
+chr2	82815451	82816236
+chr2	82816261	82816647
+chr2	82818378	82818748
+chr2	82820800	82821005
+chr2	85068666	85069046
+chr2	87824709	87825530
+chr2	89272789	89273133
+chr2	89827607	89827706
+chr2	89828636	89828710
+chr2	89828842	89828942
+chr2	89833685	89833793
+chr2	89839592	89839709
+chr2	89909317	89909789
+chr2	90379778	90402456
+chr2	92081223	92081398
+chr2	92188125	94293463
+chr2	94499181	94570956
+chr2	94898976	94899645
+chr2	94900639	94900840
+chr2	94901421	94901808
+chr2	97189431	97189813
+chr2	102482582	102482962
+chr2	102505606	102505987
+chr2	110072034	110072434
+chr2	110299106	110299346
+chr2	116751234	116751614
+chr2	116752004	116752448
+chr2	116752517	116752897
+chr2	117020171	117020552
+chr2	117021107	117022152
+chr2	117022438	117024038
+chr2	117024277	117025093
+chr2	117025205	117025670
+chr2	117026130	117026512
+chr2	120211535	120212064
+chr2	120212685	120213069
+chr2	120213761	120214143
+chr2	120214590	120215370
+chr2	121220135	121220515
+chr2	124680743	124681182
+chr2	125812046	125812548
+chr2	129090774	129091154
+chr2	130272174	130272615
+chr2	130273451	130273981
+chr2	130274326	130274992
+chr2	130275174	130275744
+chr2	130276119	130276500
+chr2	130277774	130278727
+chr2	130279995	130280729
+chr2	130280827	130281440
+chr2	130557359	130557607
+chr2	130563142	130563396
+chr2	131369643	131369925
+chr2	131370949	131371562
+chr2	131371916	131372361
+chr2	131372758	131373137
+chr2	131379317	131380344
+chr2	131381592	131381973
+chr2	131382344	131382728
+chr2	131382772	131382974
+chr2	131383079	131384016
+chr2	131384051	131384621
+chr2	131384898	131385281
+chr2	131385356	131385794
+chr2	140217229	140218044
+chr2	140220209	140220840
+chr2	140220940	140221140
+chr2	140221198	140222369
+chr2	140222545	140223623
+chr2	140223647	140224297
+chr2	143088644	143089042
+chr2	143089938	143090358
+chr2	143090898	143091662
+chr2	143092255	143092646
+chr2	143093556	143093941
+chr2	143094515	143094999
+chr2	143095614	143095994
+chr2	143096048	143096428
+chr2	143096470	143097336
+chr2	143097466	143097981
+chr2	143100621	143101005
+chr2	147048574	147048955
+chr2	147244849	147245229
+chr2	147265034	147265432
+chr2	148822913	148823295
+chr2	148881545	148882032
+chr2	155196092	155196473
+chr2	155263345	155264313
+chr2	155264362	155264562
+chr2	155264599	155264982
+chr2	155311420	155311995
+chr2	155313539	155313922
+chr2	156828628	156829008
+chr2	162517271	162517651
+chr2	164117001	164117382
+chr2	166414323	166414779
+chr2	167378863	167379244
+chr2	168652433	168652813
+chr2	179739184	179739689
+chr2	190593881	190594262
+chr2	196204680	196205060
+chr2	201212170	201212612
+chr2	201212648	201212854
+chr2	201212903	201213386
+chr2	201214659	201215040
+chr2	201549404	201549784
+chr2	201550130	201550513
+chr2	201557568	201557948
+chr2	202614117	202614527
+chr2	202615371	202615757
+chr2	202617016	202617398
+chr2	202618435	202618819
+chr2	202619754	202620134
+chr2	211773627	211774158
+chr2	211774322	211775192
+chr2	211775641	211776712
+chr2	211777034	211777417
+chr2	211777802	211778269
+chr2	211778916	211779562
+chr2	215573163	215573544
+chr2	226722088	226722596
+chr2	237521663	237522775
+chr2	237522862	237523652
+chr20	5999469	5999849
+chr20	9168743	9169145
+chr20	10441916	10442296
+chr20	13167142	13167534
+chr20	18449173	18449556
+chr20	22078162	22078542
+chr20	24024376	24024757
+chr20	26438448	28554562
+chr20	28644084	29015573
+chr20	29125977	29294639
+chr20	30744370	30744939
+chr20	30746748	30747241
+chr20	31051540	31106909
+chr20	31157044	31159116
+chr20	31161652	31223331
+chr20	34688743	34689039
+chr20	47894699	47896109
+chr20	57063873	57064279
+chr20	57357555	57358134
+chr20	57358221	57359428
+chr20	57359451	57360972
+chr20	63644937	63645318
+chr21	6369257	6372342
+chr21	7201205	7327885
+chr21	7919585	7919691
+chr21	8211710	8211892
+chr21	8212412	8212570
+chr21	8213694	8213987
+chr21	8219372	8220330
+chr21	8234456	8234568
+chr21	8394767	8394902
+chr21	8395471	8395591
+chr21	8396751	8397011
+chr21	8445918	8446080
+chr21	8446629	8446729
+chr21	8446925	8447070
+chr21	8595669	8595768
+chr21	8844362	8844855
+chr21	8846669	8847382
+chr21	10014674	10015194
+chr21	10650900	12965800
+chr21	16645305	16645685
+chr21	32095835	32096215
+chr21	35890413	35890796
+chr21	44474913	44475301
+chr21	45376056	45376517
+chr22	10863370	10863448
+chr22	11210951	11215489
+chr22	11854150	11854643
+chr22	11856460	11857173
+chr22	11974159	11974336
+chr22	12135181	12135894
+chr22	12137711	12138204
+chr22	12691742	12694097
+chr22	12954427	15057495
+chr22	15153934	15211502
+chr22	15940533	16085728
+chr22	32894952	32895345
+chr22	33819338	33819538
+chr22	35885491	35885898
+chr22	36172705	36173085
+chr22	36177875	36178257
+chr22	46470112	46470493
+chr22	50086003	50086529
+chr22	50806858	50808224
+chr3	3571912	3572292
+chr3	24705149	24705529
+chr3	25467328	25467722
+chr3	29797534	29797914
+chr3	33548103	33548483
+chr3	40252107	40253916
+chr3	41532177	41532556
+chr3	43229296	43229733
+chr3	68658875	68659467
+chr3	68670345	68670734
+chr3	73054640	73055020
+chr3	82655447	82655827
+chr3	89588895	89589538
+chr3	90269605	90722189
+chr3	90774880	91249595
+chr3	91519649	93657524
+chr3	93705477	93800019
+chr3	96475262	96475643
+chr3	96617014	96618680
+chr3	106894019	106894441
+chr3	106895181	106895568
+chr3	106896124	106896504
+chr3	106898661	106899022
+chr3	106899753	106900122
+chr3	106901799	106902741
+chr3	106903188	106903605
+chr3	119947198	119947578
+chr3	120721858	120722610
+chr3	122688557	122688938
+chr3	125982519	125982900
+chr3	127005357	127005745
+chr3	128988979	128989359
+chr3	137095968	137096348
+chr3	142662232	142662612
+chr3	152919604	152919995
+chr3	153658704	153659087
+chr3	160947473	160948127
+chr3	166159726	166160108
+chr3	166160260	166160644
+chr3	166161631	166162087
+chr3	166226563	166226945
+chr3	166232406	166232886
+chr3	166232970	166233355
+chr3	166474023	166474223
+chr3	171534313	171534700
+chr3	177010776	177011156
+chr3	192880587	192880967
+chr4	5404508	5404897
+chr4	12640142	12640815
+chr4	14506099	14506467
+chr4	17061824	17062213
+chr4	18949310	18949691
+chr4	22502173	22502553
+chr4	25717756	25718136
+chr4	25718275	25718655
+chr4	25719398	25719626
+chr4	27730251	27730747
+chr4	30884524	30884906
+chr4	32280109	32280489
+chr4	41023064	41023448
+chr4	47772100	47772544
+chr4	49136056	49136102
+chr4	49141052	49141147
+chr4	49246355	49246848
+chr4	49548607	49549100
+chr4	49631231	49658125
+chr4	49708086	51743949
+chr4	51793952	51817249
+chr4	55327979	55328462
+chr4	64606369	64606752
+chr4	64606841	64607360
+chr4	64607395	64607789
+chr4	64607976	64608801
+chr4	64608937	64609326
+chr4	64609811	64610876
+chr4	64611176	64611617
+chr4	66065193	66065631
+chr4	68050141	68050521
+chr4	68572333	68572774
+chr4	78008402	78008882
+chr4	83383282	83383662
+chr4	89731703	89732163
+chr4	92701787	92702300
+chr4	107501924	107502304
+chr4	112372589	112372969
+chr4	116296652	116297040
+chr4	116297165	116297545
+chr4	116297659	116298726
+chr4	116299003	116300416
+chr4	128081280	128081956
+chr4	140929567	140929947
+chr4	143017907	143018107
+chr4	143347973	143348354
+chr4	144379497	144379877
+chr4	155076906	155077288
+chr4	155452733	155452935
+chr4	155453928	155454313
+chr4	155454407	155455447
+chr4	155455566	155455766
+chr4	155457624	155458008
+chr4	155459547	155459747
+chr4	155460171	155460553
+chr4	155461093	155461689
+chr4	155462078	155463456
+chr4	155463701	155464839
+chr4	155464895	155465305
+chr4	155465580	155466624
+chr4	157628391	157628774
+chr4	160044429	160044815
+chr4	161449477	161449857
+chr4	161788291	161788671
+chr4	162421207	162421721
+chr4	172036714	172037094
+chr4	179069259	179069639
+chr4	183489243	183489623
+chr4	189844495	189844576
+chr5	12284	12523
+chr5	12952	13361
+chr5	5395563	5395943
+chr5	5396182	5396616
+chr5	5396675	5397057
+chr5	8619083	8619464
+chr5	8619927	8620307
+chr5	8620707	8621192
+chr5	8621953	8622333
+chr5	8622354	8622753
+chr5	32927394	32927776
+chr5	37164286	37164673
+chr5	45913363	50265419
+chr5	60761358	60762176
+chr5	66253509	66253889
+chr5	73775720	73776112
+chr5	79089860	79090240
+chr5	80649841	80652548
+chr5	94567275	94571098
+chr5	97678633	97679016
+chr5	98409947	98410327
+chr5	98410700	98411257
+chr5	99813005	99813388
+chr5	100045805	100055225
+chr5	106553187	106553689
+chr5	111488864	111489244
+chr5	119127218	119127602
+chr5	121030820	121031445
+chr5	122338658	122339042
+chr5	123760111	123760622
+chr5	123760719	123761918
+chr5	134923133	134928692
+chr5	136533606	136533986
+chr5	137305006	137305387
+chr5	152198765	152199145
+chr5	160600365	160600745
+chr5	163146853	163147234
+chr5	163959711	163960091
+chr5	164673914	164674288
+chr5	166530241	166530641
+chr5	170635389	170635774
+chr6	1705930	1706304
+chr6	3943769	3944149
+chr6	29454054	29454435
+chr6	32706020	32706850
+chr6	43490986	43491370
+chr6	54899048	54899248
+chr6	58554346	59830578
+chr6	61278527	61521106
+chr6	61573960	61574809
+chr6	72747981	72748361
+chr6	72799169	72799549
+chr6	76708390	76708770
+chr6	88555202	88555591
+chr6	91726616	91727363
+chr6	94446937	94447370
+chr6	96941571	96941951
+chr6	104699855	104700055
+chr6	114377334	114377534
+chr6	122764824	122765204
+chr6	126478329	126478709
+chr6	127735330	127735710
+chr6	132799554	132799939
+chr6	133150492	133150881
+chr6	133930809	133931190
+chr6	138133082	138133462
+chr6	143077647	143078031
+chr6	153666229	153666618
+chr6	153667363	153667744
+chr6	153668187	153668753
+chr6	153669025	153669419
+chr6	156547729	156548118
+chr6	163638068	163638448
+chr7	18021726	18022106
+chr7	22748471	22748854
+chr7	33749120	33749500
+chr7	36228567	36229008
+chr7	37387570	37387950
+chr7	45251808	45252289
+chr7	55369049	55369429
+chr7	57167688	57168071
+chr7	57168472	57168852
+chr7	57169046	57169430
+chr7	57169550	57169932
+chr7	57170307	57170523
+chr7	57170675	57171410
+chr7	57171502	57172122
+chr7	57173798	57174181
+chr7	57174854	57175239
+chr7	57185615	57185995
+chr7	57186105	57186589
+chr7	57187287	57188033
+chr7	57188305	57188872
+chr7	57189116	57189730
+chr7	57190949	57191332
+chr7	57191618	57191818
+chr7	57192132	57192860
+chr7	57193489	57193872
+chr7	57193974	57194701
+chr7	57194829	57195210
+chr7	57196302	57197490
+chr7	57198263	57198644
+chr7	57879605	58032504
+chr7	58166363	62995324
+chr7	63094673	63095057
+chr7	64104133	64104513
+chr7	64105294	64106415
+chr7	64106627	64107010
+chr7	64108329	64108798
+chr7	64110007	64110707
+chr7	64111376	64111804
+chr7	64111957	64112849
+chr7	67627830	67628213
+chr7	68097607	68097990
+chr7	68736347	68736811
+chr7	69331805	69332005
+chr7	69332037	69332438
+chr7	69333013	69333393
+chr7	69333597	69334167
+chr7	72088575	72088955
+chr7	83100026	83100406
+chr7	83469984	83470184
+chr7	83855080	83855464
+chr7	95851249	95851629
+chr7	104989516	104989896
+chr7	112372484	112372865
+chr7	112374724	112374950
+chr7	117263552	117264184
+chr7	117264231	117264614
+chr7	130116678	130117058
+chr7	141173000	141173384
+chr7	141801916	141802451
+chr7	141802901	141803366
+chr7	141804074	141804274
+chr7	141804814	141805507
+chr7	142665099	142667846
+chr7	143187483	143187863
+chr7	145997159	145997608
+chr7	150131843	150132229
+chr7	153968598	153968979
+chr7	159294463	159294846
+chr8	13353292	13353679
+chr8	16056863	16057063
+chr8	18849121	18849571
+chr8	20551162	20551554
+chr8	32805708	32806092
+chr8	33010514	33010894
+chr8	33011359	33014071
+chr8	33014510	33014895
+chr8	33015020	33015853
+chr8	36277446	36278060
+chr8	36278272	36278791
+chr8	36278835	36279634
+chr8	40070431	40070867
+chr8	43237631	43242390
+chr8	43937900	45969600
+chr8	46827305	46827914
+chr8	46828298	46829961
+chr8	46830195	46831222
+chr8	46837581	46837961
+chr8	46838101	46838484
+chr8	50758259	50758639
+chr8	56736733	56736933
+chr8	61303079	61303460
+chr8	67580689	67581493
+chr8	67581588	67581972
+chr8	67582178	67582568
+chr8	67585216	67585693
+chr8	67585787	67586175
+chr8	67587282	67587922
+chr8	69102851	69103234
+chr8	72985528	72985923
+chr8	74828644	74829025
+chr8	76201592	76202319
+chr8	76645407	76645800
+chr8	97907908	97908279
+chr8	99495689	99496133
+chr8	102774315	102774695
+chr8	103082925	103083379
+chr8	103083704	103084399
+chr8	103084730	103085110
+chr8	103085323	103085806
+chr8	103086859	103087242
+chr8	108533901	108534281
+chr8	110933150	110933533
+chr8	110934510	110935010
+chr8	111248936	111249316
+chr8	120224204	120224584
+chr8	127053876	127054257
+chr8	127968653	127969034
+chr8	133615761	133616142
+chr8	133755390	133755856
+chr9	5091131	5091511
+chr9	5091962	5093013
+chr9	5093063	5094123
+chr9	5094192	5094697
+chr9	5094931	5095816
+chr9	5096206	5096816
+chr9	5097188	5097890
+chr9	5098134	5098516
+chr9	5099352	5099552
+chr9	5100044	5100427
+chr9	5108063	5108592
+chr9	5109193	5109986
+chr9	5110030	5110411
+chr9	9896970	9897350
+chr9	15866612	15866992
+chr9	18336471	18336854
+chr9	31498260	31498640
+chr9	33656533	33658316
+chr9	33658346	33659299
+chr9	34998988	34999474
+chr9	36466192	36466572
+chr9	43153721	45525161
+chr9	64045550	64046043
+chr9	64047855	64048422
+chr9	65048153	65079624
+chr9	68251002	68251071
+chr9	72788174	72788555
+chr9	78741395	78741775
+chr9	78742155	78742969
+chr9	78743199	78743630
+chr9	78744108	78744492
+chr9	78810721	78811113
+chr9	79804550	79804933
+chr9	80564643	80565085
+chr9	80565478	80565941
+chr9	81747641	81748021
+chr9	82427689	82428071
+chr9	92108965	92109347
+chr9	92539106	92539763
+chr9	95876956	95877338
+chr9	117109914	117110296
+chr9	122505687	122506067
+chr9	129878699	129879081
+chr9	134164478	134165354
+chr9	134170819	134171060
+chrX	4059512	4059712
+chrX	5168678	5169232
+chrX	5169733	5170646
+chrX	15727702	15728089
+chrX	17116414	17116794
+chrX	24056083	24056470
+chrX	24375345	24375545
+chrX	33762401	33762781
+chrX	55178596	55179289
+chrX	55179434	55180459
+chrX	55181196	55182790
+chrX	55183051	55184112
+chrX	58061543	62821716
+chrX	62841379	62841765
+chrX	62842257	62842639
+chrX	70119464	70119845
+chrX	70127233	70127620
+chrX	77501934	77502314
+chrX	78561721	78561921
+chrX	84403779	84404168
+chrX	100027094	100027475
+chrX	102010329	102010712
+chrX	102011531	102011915
+chrX	102772405	102772791
+chrX	102785904	102786287
+chrX	102798001	102798386
+chrX	102802747	102803161
+chrX	102809395	102809788
+chrX	104409869	104410249
+chrX	106239694	106239894
+chrX	111416893	111417294
+chrX	126471558	126473451
+chrX	126728884	126729272
+chrX	126729326	126729709
+chrX	126729837	126730217
+chrX	126730716	126731106
+chrX	126731624	126732029
+chrX	129983338	129983538
+chrX	133041871	133042251
+chrX	135292293	135292493
+chrX	143430213	143430837
+chrX	143431144	143431537
+chrX	143431716	143432219
+chrX	143432410	143433212
+chrX	143433510	143434156
+chrX	143543636	143544023
+chrX	146995842	146996224
+chrY	4344757	4344879
+chrY	9141870	9141995
+chrY	10203380	10266932
+chrY	10316749	10544446
+chrY	10594583	10626838
+chrY	10663669	10663716
+chrY	10744417	10921497
+chrY	11290797	11334278
+chrY	11493053	11592850
+chrY	11671014	11671046
+chrY	11721528	11749472
+chrY	56694632	56889743

data/chrom_size_hg38.txt

@@ -0,0 +1,23 @@
+chr1	248956422
+chr2	242193529
+chr3	198295559
+chr4	190214555
+chr5	181538259
+chr6	170805979
+chr7	159345973
+chrX	156040895
+chr8	145138636
+chr9	138394717
+chr11	135086622
+chr10	133797422
+chr12	133275309
+chr13	114364328
+chr14	107043718
+chr15	101991189
+chr16	90338345
+chr17	83257441
+chr18	80373285
+chr20	64444167
+chr19	58617616
+chr22	50818468
+chr21	46709983

data/epigenomes.txt

@@ -0,0 +1,245 @@
+CTCF
+RAD21
+ATF7
+CEBPB
+MAX
+CREB1
+DPF2
+POLR2A
+JUND
+ZBTB33
+ELF1
+MAZ
+MAFK
+REST
+EGR1
+EP300
+SIN3A
+ZBTB40
+GABPA
+RCOR1
+ESRRA
+NBN
+NRF1
+HDAC2
+MYC
+ARID3A
+FOXM1
+RAD51
+MTA3
+TCF12
+HCFC1
+FOXK2
+SRF
+SP1
+ELK1
+RFX5
+SUZ12
+ZNF687
+PKNOX1
+IKZF1
+NFIC
+CEBPG
+ZFX
+NR2F1
+ATF2
+BHLHE40
+SMC3
+POLR2AphosphoS5
+ZNF217
+NR2C2
+ZNF143
+NR2F2
+CREM
+MLLT1
+TRIM22
+TEAD4
+YY1
+MTA2
+KDM1A
+MNT
+TCF3
+UBTF
+MXI1
+ZNF384
+RFX1
+ZNF592
+TCF7L2
+SMARCA5
+PML
+E2F1
+SMARCE1
+MTA1
+TAF1
+FOXA1
+CBFB
+GATAD2B
+ZNF24
+TOE1
+CBX5
+JUN
+TARDBP
+ZFP36
+NFATC3
+ETV6
+MEF2A
+ARNT
+TCF7
+CHD2
+HDGF
+ETS1
+TBL1XR1
+SMAD1
+CUX1
+POLR2AphosphoS2
+USF2
+ZSCAN29
+E2F4
+LARP7
+BMI1
+ZKSCAN1
+CHD1
+BCLAF1
+CHD4
+NFRKB
+E2F8
+GTF2F1
+ZNF207
+SREBF1
+BRCA1
+NONO
+ZZZ3
+NFXL1
+E4F1
+ZMYM3
+HDAC1
+SMAD4
+MAFF
+CTBP1
+ATF3
+SPI1
+GABPB1
+FOSL2
+MEIS2
+FOSL1
+STAG1
+GATAD2A
+ZNF281
+FOXA3
+PHF8
+EED
+KDM5B
+KLF16
+MAFG
+ETV5
+TEAD1
+PRDM10
+TFE3
+TFDP1
+PHF21A
+SKIL
+PATZ1
+ATF4
+BCL6
+ZNF766
+ERF
+NR2F6
+ZNF574
+IRF2
+BCL3
+GATA2
+JUNB
+BACH1
+RB1
+ZNF644
+ZBTB7A
+MEF2D
+KAT8
+ZBTB7B
+SOX6
+NKRF
+RXRA
+NFE2
+RBPJ
+ADNP
+GMEB1
+NCOR1
+RREB1
+ZKSCAN8
+PHF20
+SMAD5
+ZNF512
+ELF4
+NFYB
+HOMEZ
+EHMT2
+MYBL2
+THRA
+RNF2
+SMAD3
+ZNF444
+E2F5
+KLF9
+HMG20A
+ZNF12
+NR2C1
+MBD2
+KLF6
+NFIB
+NCOA1
+CBX1
+BRD4
+RBBP5
+FOXJ3
+ZNF7
+HIVEP1
+ZNF274
+ZEB1
+STAT5A
+HINFP
+TBP
+ZBED1
+ZHX2
+NEUROD1
+ZNF740
+ZNF83
+NFYA
+ZNF280B
+IRF1
+CREB3
+MLX
+ATF1
+ZNF318
+SMARCC2
+ZFP1
+SIX4
+ZNF512B
+PRDM15
+CLOCK
+HMG20B
+SIN3B
+KLF13
+ZNF507
+IRF5
+ATF6
+RELA
+ZNF639
+ZC3H8
+HES1
+KDM4B
+DNMT1
+DDX20
+ZC3H4
+NFE2L1
+ZNF282
+PBX2
+ZNF3
+H3K27me3
+H3K9ac
+H3K4me1
+H3K27ac
+H3K36me3
+H2AFZ
+H3K4me3
+H3K9me3
+H3K4me2

models/cage.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bd3a532d752b3067f5424a0a0357abc9eba010f878c5a324177aedbc80909d12
+size 128747867

models/epi_bind.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6d71047f59dab64333b93d2cb10640155ca9749ebf8856bb2dacb79b6c4b9ae5
+size 59823081

models/epi_track.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bb75d025e1e689165b56a46209ef9c9e9c5f575d7751886a83c1b470d9a69b5b
+size 128348451

models/hic.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cbfe0c2e5b9eb0a4e2f2eba7076e0263f42d4e763b04e0368aabf25d46d6b8a9
+size 54561869

models/microc.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8244fc07de18d110e28820b860bf8715a87a58b1835615b8253ffc43727aacc1
+size 140192327

pretrain/layers.py

@@ -0,0 +1,87 @@
+import torch
+import torch.nn as nn
+import numpy as np
+import torch.nn.functional as F
+
+
+class CNN(nn.Module):
+    def __init__(self):
+        super(CNN, self).__init__()
+        conv_kernel_size1 = 10
+        conv_kernel_size2 = 8
+        pool_kernel_size1 = 5
+        pool_kernel_size2 = 4
+        self.conv_net = nn.Sequential(
+            nn.Conv1d(5, 256, kernel_size=conv_kernel_size1),
+            nn.ReLU(inplace=True),
+            nn.Dropout(p=0.1),
+            nn.Conv1d(256, 256, kernel_size=conv_kernel_size1),
+            nn.BatchNorm1d(256),
+            nn.ReLU(inplace=True),
+            nn.MaxPool1d(kernel_size=pool_kernel_size1, stride=pool_kernel_size1),
+            nn.Dropout(p=0.1),
+            nn.Conv1d(256, 360, kernel_size=conv_kernel_size2),
+            nn.ReLU(inplace=True),
+            nn.Dropout(p=0.1),
+            nn.Conv1d(360, 360, kernel_size=conv_kernel_size2),
+            nn.BatchNorm1d(360),
+            nn.ReLU(inplace=True),
+            nn.MaxPool1d(kernel_size=pool_kernel_size2, stride=pool_kernel_size2),
+            nn.Dropout(p=0.1),
+            nn.Conv1d(360, 512, kernel_size=conv_kernel_size2),
+            nn.ReLU(inplace=True),
+            nn.Dropout(p=0.2),
+            nn.Conv1d(512, 512, kernel_size=conv_kernel_size2),
+            nn.BatchNorm1d(512),
+            nn.ReLU(inplace=True),
+            nn.Dropout(p=0.2))
+        self.num_channels = 512
+    def forward(self, x):
+        out = self.conv_net(x)
+        return out
+
+
+# borrow from https://github.com/Alibaba-MIIL/ASL/blob/main/src/loss_functions/losses.py
+class Balanced_AsymmetricLoss(nn.Module):
+    def __init__(self, gamma_neg=4, gamma_pos=1, clip=0.05, alpha=None, eps=1e-8, disable_torch_grad_focal_loss=True):
+        super(Balanced_AsymmetricLoss, self).__init__()
+
+        self.gamma_neg = gamma_neg
+        self.gamma_pos = gamma_pos
+        self.clip = clip
+        self.disable_torch_grad_focal_loss = disable_torch_grad_focal_loss
+        self.eps = eps
+        self.alpha = alpha
+
+    def forward(self, x, y, mask):
+        # Calculating Probabilities
+        assert y.shape == mask.shape
+        x_sigmoid = torch.sigmoid(x)
+        xs_pos = x_sigmoid
+        xs_neg = 1 - x_sigmoid
+        # Asymmetric Clipping
+        if self.clip is not None and self.clip > 0:
+            xs_neg = (xs_neg + self.clip).clamp(max=1)
+
+        # Basic CE calculation
+        los_pos = y * torch.log(xs_pos.clamp(min=self.eps))
+        los_neg = (1 - y) * torch.log(xs_neg.clamp(min=self.eps))
+        if self.alpha is not None:
+            los_pos = self.alpha * los_pos
+        loss = los_pos + los_neg
+        # Asymmetric Focusing
+        if self.gamma_neg > 0 or self.gamma_pos > 0:
+            if self.disable_torch_grad_focal_loss:
+                torch.set_grad_enabled(False)
+            pt0 = xs_pos * y
+            pt1 = xs_neg * (1 - y)  # pt = p if t > 0 else 1-p
+            pt = pt0 + pt1
+            one_sided_gamma = self.gamma_pos * y + self.gamma_neg * (1 - y)
+            one_sided_w = torch.pow(1 - pt, one_sided_gamma)
+            if self.disable_torch_grad_focal_loss:
+                torch.set_grad_enabled(True)
+            loss *= one_sided_w
+        loss *= mask
+        return -loss.sum() / (torch.sum(mask) + self.eps)
+
+

pretrain/model.py

@@ -0,0 +1,72 @@
+#most of the codes below are copied from Query2label
+import torch,math
+import numpy as np
+from torch import nn, Tensor
+from pretrain.layers import CNN
+from pretrain.transformer import Transformer
+
+class GroupWiseLinear(nn.Module):
+    def __init__(self, num_class, hidden_dim, bias=True):
+        super().__init__()
+        self.num_class = num_class
+        self.hidden_dim = hidden_dim
+        self.bias = bias
+        self.W = nn.Parameter(torch.Tensor(1, num_class, hidden_dim))
+        if bias:
+            self.b = nn.Parameter(torch.Tensor(1, num_class))
+        self.reset_parameters()
+    def reset_parameters(self):
+        stdv = 1. / math.sqrt(self.W.size(2))
+        for i in range(self.num_class):
+            self.W[0][i].data.uniform_(-stdv, stdv)
+        if self.bias:
+            for i in range(self.num_class):
+                self.b[0][i].data.uniform_(-stdv, stdv)
+    def forward(self, x):
+        x = (self.W * x).sum(-1)
+        if self.bias:
+            x = x + self.b
+        return x
+
+class Tranmodel(nn.Module):
+    def __init__(self, backbone, transfomer, num_class):
+        super().__init__()
+        self.backbone = backbone
+        self.transformer = transfomer
+        self.num_class = num_class
+        hidden_dim = transfomer.d_model
+        self.label_input = torch.Tensor(np.arange(num_class)).view(1, -1).long()
+        self.input_proj = nn.Conv1d(backbone.num_channels, hidden_dim, kernel_size=1)
+        self.query_embed = nn.Embedding(num_class, hidden_dim)
+        self.fc = GroupWiseLinear(num_class, hidden_dim, bias=True)
+
+    def forward(self, input):
+        src = self.backbone(input)
+        label_inputs=self.label_input.repeat(src.size(0),1).to(input.device)
+        label_embed=self.query_embed(label_inputs)
+        src=self.input_proj(src)
+        hs = self.transformer(src, label_embed)
+        out = self.fc(hs)
+        return out
+
+def build_backbone():
+    model = CNN()
+    return model
+def build_transformer(args):
+    return Transformer(
+        d_model=args.hidden_dim,
+        dropout=args.dropout,
+        nhead=args.nheads,
+        dim_feedforward=args.dim_feedforward,
+        num_encoder_layers=args.enc_layers,
+        num_decoder_layers=args.dec_layers
+    )
+def build_epd_model(args):
+    backbone = build_backbone()
+    transformer = build_transformer(args)
+    model = Tranmodel(
+            backbone=backbone,
+            transfomer=transformer,
+            num_class=args.num_class,
+        )
+    return model

pretrain/track/layers.py

@@ -0,0 +1,299 @@
+# some of following codes are borrowed from https://github.com/lucidrains/enformer-pytorch
+
+import math
+import torch
+from torch import nn, einsum
+from einops import rearrange, reduce
+from einops.layers.torch import Rearrange
+import torch.nn.functional as F
+
+def exists(val):
+    return val is not None
+
+def default(val, d):
+    return val if exists(val) else d
+
+def map_values(fn, d):
+    return {key: fn(values) for key, values in d.items()}
+
+def exponential_linspace_int(start, end, num, divisible_by = 1):
+    def _round(x):
+        return int(round(x / divisible_by) * divisible_by)
+
+    base = math.exp(math.log(end / start) / (num - 1))
+    return [_round(start * base**i) for i in range(num)]
+
+def log(t, eps = 1e-20):
+    return torch.log(t.clamp(min = eps))
+
+# losses and metrics
+
+
+def pearson_corr_coef(x, y, eps = 1e-8):
+    x2 = x * x
+    y2 = y * y
+    xy = x * y
+    ex = x.mean(dim = 1)
+    ey = y.mean(dim = 1)
+    exy = xy.mean(dim = 1)
+    ex2 = x2.mean(dim = 1)
+    ey2 = y2.mean(dim = 1)
+    r = (exy - ex * ey) / (torch.sqrt(ex2 - (ex * ex)) * torch.sqrt(ey2 - (ey * ey)) + eps)
+    return r.mean(dim = -1)
+
+# relative positional encoding functions
+
+def get_positional_features_exponential(positions, features, seq_len, min_half_life = 3.):
+    max_range = math.log(seq_len) / math.log(2.)
+    half_life = 2 ** torch.linspace(min_half_life, max_range, features, device = positions.device)
+    half_life = half_life[None, ...]
+    positions = positions.abs()[..., None]
+    return torch.exp(-math.log(2.) / half_life * positions)
+
+def get_positional_features_central_mask(positions, features, seq_len):
+    center_widths = 2 ** torch.arange(1, features + 1, device = positions.device).float()
+    center_widths = center_widths - 1
+    return (center_widths[None, ...] > positions.abs()[..., None]).float()
+
+def gamma_pdf(x, concentration, rate):
+    log_unnormalized_prob = torch.xlogy(concentration - 1., x) - rate * x
+    log_normalization = (torch.lgamma(concentration) - concentration * torch.log(rate))
+    return torch.exp(log_unnormalized_prob - log_normalization)
+
+def get_positional_features_gamma(positions, features, seq_len, stddev = None, start_mean = None, eps = 1e-8):
+    if not exists(stddev):
+        stddev = seq_len / (2 * features)
+
+    if not exists(start_mean):
+        start_mean = seq_len / features
+
+    mean = torch.linspace(start_mean, seq_len, features, device = positions.device)
+    mean = mean[None, ...]
+    concentration = (mean / stddev) ** 2
+    rate = mean / stddev ** 2
+    probabilities = gamma_pdf(positions.float().abs()[..., None], concentration, rate)
+    probabilities = probabilities + eps
+    outputs = probabilities / torch.amax(probabilities)
+    return outputs
+
+def get_positional_embed(seq_len, feature_size, device):
+    distances = torch.arange(-seq_len + 1, seq_len, device = device)
+
+    feature_functions = [
+        get_positional_features_exponential,
+        get_positional_features_central_mask,
+        get_positional_features_gamma
+    ]
+
+    num_components = len(feature_functions) * 2
+
+    if (feature_size % num_components) != 0:
+        raise ValueError(f'feature size is not divisible by number of components ({num_components})')
+
+    num_basis_per_class = feature_size // num_components
+
+    embeddings = []
+    for fn in feature_functions:
+        embeddings.append(fn(distances, num_basis_per_class, seq_len))
+
+    embeddings = torch.cat(embeddings, dim = -1)
+    embeddings = torch.cat((embeddings, torch.sign(distances)[..., None] * embeddings), dim = -1)
+    return embeddings
+
+def relative_shift(x):
+    to_pad = torch.zeros_like(x[..., :1])
+    x = torch.cat((to_pad, x), dim = -1)
+    _, h, t1, t2 = x.shape
+    x = x.reshape(-1, h, t2, t1)
+    x = x[:, :, 1:, :]
+    x = x.reshape(-1, h, t1, t2 - 1)
+    return x[..., :((t2 + 1) // 2)]
+
+# classes
+
+class Residual(nn.Module):
+    def __init__(self, fn):
+        super().__init__()
+        self.fn = fn
+
+    def forward(self, x, **kwargs):
+        return self.fn(x, **kwargs) + x
+
+class GELU(nn.Module):
+    def forward(self, x):
+        return torch.sigmoid(1.702 * x) * x
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,dim,num_rel_pos_features,
+        heads = 8,
+        dim_key = 64,
+        dim_value = 64,
+        dropout = 0.,
+        pos_dropout = 0.
+    ):
+        super().__init__()
+        self.scale = dim_key ** -0.5
+        self.heads = heads
+
+        self.to_q = nn.Linear(dim, dim_key * heads, bias = False)
+        self.to_k = nn.Linear(dim, dim_key * heads, bias = False)
+        self.to_v = nn.Linear(dim, dim_value * heads, bias = False)
+
+        self.to_out = nn.Linear(dim_value * heads, dim)
+        nn.init.zeros_(self.to_out.weight)
+        nn.init.zeros_(self.to_out.bias)
+
+        # relative positional encoding
+
+        self.num_rel_pos_features = num_rel_pos_features
+
+        self.to_rel_k = nn.Linear(num_rel_pos_features, dim_key * heads, bias = False)
+        self.rel_content_bias = nn.Parameter(torch.randn(1, heads, 1, dim_key))
+        self.rel_pos_bias = nn.Parameter(torch.randn(1, heads, 1, dim_key))
+
+        self.pos_dropout = nn.Dropout(pos_dropout)
+        self.attn_dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        n, h, device = x.shape[-2], self.heads, x.device
+
+        q = self.to_q(x)
+        k = self.to_k(x)
+        v = self.to_v(x)
+
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), (q, k, v))
+
+        q = q * self.scale
+
+        content_logits = einsum('b h i d, b h j d -> b h i j', q + self.rel_content_bias, k)
+
+        positions = get_positional_embed(n, self.num_rel_pos_features, device)
+        positions = self.pos_dropout(positions)
+        rel_k = self.to_rel_k(positions)
+
+        rel_k = rearrange(rel_k, 'n (h d) -> h n d', h = h)
+        rel_logits = einsum('b h i d, h j d -> b h i j', q + self.rel_pos_bias, rel_k)
+        rel_logits = relative_shift(rel_logits)
+
+        logits = content_logits + rel_logits
+        attn = logits.softmax(dim = -1)
+        attn = self.attn_dropout(attn)
+
+        out = einsum('b h i j, b h j d -> b h i d', attn, v)
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        return self.to_out(out)
+
+class Enformer(nn.Module):
+    def __init__(
+        self,
+        dim = 512,
+        depth = 4,
+        heads = 6,
+        attn_dim_key = 64,
+        dropout_rate = 0.2,
+        attn_dropout = 0.05,
+        pos_dropout = 0.01,
+    ):
+        super().__init__()
+        self.dim = dim
+        transformer = []
+        for _ in range(depth):
+            transformer.append(nn.Sequential(
+                Residual(nn.Sequential(
+                    nn.LayerNorm(dim),
+                    Attention(
+                        dim,
+                        heads = heads,
+                        dim_key = attn_dim_key,
+                        dim_value = dim // heads,
+                        dropout = attn_dropout,
+                        pos_dropout = pos_dropout,
+                        num_rel_pos_features = dim // heads
+                    ),
+                    nn.Dropout(dropout_rate)
+                )),
+                Residual(nn.Sequential(
+                    nn.LayerNorm(dim),
+                    nn.Linear(dim, dim * 2),
+                    nn.Dropout(dropout_rate),
+                    nn.ReLU(),
+                    nn.Linear(dim * 2, dim),
+                    nn.Dropout(dropout_rate)
+                ))
+            ))
+
+        self.transformer = nn.Sequential(
+            # Rearrange('b d n -> b n d'),
+            *transformer
+        )
+
+    def forward(self,x):
+
+        x = self.transformer(x)
+
+        return x
+
+class CNN(nn.Module):
+    def __init__(self):
+        super(CNN, self).__init__()
+        conv_kernel_size1 = 10
+        conv_kernel_size2 = 8
+        pool_kernel_size1 = 5
+        pool_kernel_size2 = 4
+        self.conv_net = nn.Sequential(
+            nn.Conv1d(5, 256, kernel_size=conv_kernel_size1),
+            nn.ReLU(inplace=True),
+            nn.Dropout(p=0.1),
+            nn.Conv1d(256, 256, kernel_size=conv_kernel_size1),
+            # nn.GroupNorm(16, 256),
+            nn.BatchNorm1d(256,track_running_stats=False),
+            nn.ReLU(inplace=True),
+            nn.MaxPool1d(kernel_size=pool_kernel_size1, stride=pool_kernel_size1),
+            nn.Dropout(p=0.1),
+            nn.Conv1d(256, 360, kernel_size=conv_kernel_size2),
+            nn.ReLU(inplace=True),
+            nn.Dropout(p=0.1),
+            nn.Conv1d(360, 360, kernel_size=conv_kernel_size2),
+            nn.BatchNorm1d(360,track_running_stats=False),
+            # nn.GroupNorm(36, 360),
+            nn.ReLU(inplace=True),
+            nn.MaxPool1d(kernel_size=pool_kernel_size2, stride=pool_kernel_size2),
+            nn.Dropout(p=0.1),
+            nn.Conv1d(360, 512, kernel_size=conv_kernel_size2),
+            nn.ReLU(inplace=True),
+            nn.Dropout(p=0.2),
+            nn.Conv1d(512, 512, kernel_size=conv_kernel_size2),
+            nn.BatchNorm1d(512,track_running_stats=False),
+            # nn.GroupNorm(32, 512),
+            nn.ReLU(inplace=True),
+            nn.Dropout(p=0.2))
+        self.num_channels = 512
+    def forward(self, x):
+        out = self.conv_net(x)
+        return out
+
+
+class AttentionPool(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.pool_fn = Rearrange('b (n p) d-> b n p d', n=1)
+        self.to_attn_logits = nn.Parameter(torch.eye(dim))
+
+    def forward(self, x):
+        attn_logits = einsum('b n d, d e -> b n e', x, self.to_attn_logits)
+        x = self.pool_fn(x)
+        logits = self.pool_fn(attn_logits)
+
+        attn = logits.softmax(dim = -2)
+        return (x * attn).sum(dim = -2).squeeze()
+
+
+
+# class MSE_loss(nn.Module):
+#     def __init__(self):
+#         super().__init__()
+#     def forward(self,pred,target):
+
+

pretrain/track/model.py

@@ -0,0 +1,105 @@
+import torch,math
+from torch import nn, Tensor
+from pretrain.track.transformers import Transformer
+from einops import rearrange,repeat
+from pretrain.track.layers import CNN,Enformer,AttentionPool
+from einops.layers.torch import Rearrange
+
+import numpy as np
+import torch.nn.functional as F
+class Tranmodel(nn.Module):
+    def __init__(self, backbone, transfomer):
+        super().__init__()
+        self.backbone = backbone
+        self.transformer = transfomer
+        hidden_dim = transfomer.d_model
+        self.input_proj = nn.Conv1d(backbone.num_channels, hidden_dim, kernel_size=1)
+    def forward(self, input):
+        input=rearrange(input,'b n c l -> (b n) c l')
+        src = self.backbone(input)
+        src=self.input_proj(src)
+        src = self.transformer(src)
+        return src
+
+class PositionalEncoding(nn.Module):
+
+    def __init__(self, d_model: int, dropout: float = 0.1, max_len: int = 5000):
+        super().__init__()
+        self.dropout = nn.Dropout(p=dropout)
+
+        position = torch.arange(max_len).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
+        pe = torch.zeros(max_len, 1, d_model)
+        pe[:, 0, 0::2] = torch.sin(position * div_term)
+        pe[:, 0, 1::2] = torch.cos(position * div_term)
+        self.register_buffer('pe', pe)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = x + self.pe[:x.size(0)]
+        return self.dropout(x)
+
+class finetunemodel(nn.Module):
+    def __init__(self,pretrain_model,hidden_dim,embed_dim,bins,crop=50,num_class=245,return_embed=True):
+        super().__init__()
+        self.pretrain_model = pretrain_model
+        self.bins=bins
+        self.crop=crop
+        self.return_embed = return_embed
+        self.attention_pool = AttentionPool(hidden_dim)
+        self.project=nn.Sequential(
+            Rearrange('(b n) c -> b c n', n=bins),
+            nn.Conv1d(hidden_dim, hidden_dim, kernel_size=15, padding=7,groups=hidden_dim),
+            nn.InstanceNorm1d(hidden_dim, affine=True),
+            nn.Conv1d(hidden_dim, embed_dim, kernel_size=1),
+            nn.ReLU(inplace=True),
+            nn.Dropout(0.1),
+            nn.Conv1d(embed_dim, embed_dim, kernel_size=9, padding=4),
+            nn.InstanceNorm1d(embed_dim, affine=True),
+            nn.ReLU(inplace=True),
+            nn.Dropout(0.1),
+        )
+
+
+        self.transformer = Enformer(dim=embed_dim, depth=4, heads=8)
+        self.prediction_head=nn.Linear(embed_dim,num_class)
+
+
+    def forward(self, x):
+        x=self.pretrain_model(x)
+        x = self.attention_pool(x)
+        x = self.project(x)
+        x= rearrange(x,'b c n -> b n c')
+
+        x = self.transformer(x)
+        out = self.prediction_head(x[:, self.crop:-self.crop, :])
+        if self.return_embed:
+            return x
+        else:
+            return out
+
+
+
+
+def build_backbone():
+    model = CNN()
+    return model
+def build_transformer(args):
+    return Transformer(
+        d_model=args.hidden_dim,
+        dropout=args.dropout,
+        nhead=args.nheads,
+        dim_feedforward=args.dim_feedforward,
+        num_encoder_layers=args.enc_layers,
+        num_decoder_layers=args.dec_layers
+    )
+def build_track_model(args):
+    backbone = build_backbone()
+    transformer = build_transformer(args)
+    pretrain_model = Tranmodel(
+            backbone=backbone,
+            transfomer=transformer,
+        )
+    model=finetunemodel(pretrain_model,hidden_dim=args.hidden_dim,embed_dim=args.embed_dim,
+                        bins=args.bins,crop=args.crop,return_embed=args.return_embed)
+
+    return model

pretrain/track/transformers.py

@@ -0,0 +1,94 @@
+# most of the codes are borrowed from Query2label and DETR
+import torch
+import torch.nn.functional as F
+from torch import nn, Tensor
+from torch.nn import MultiheadAttention
+from typing import Optional
+import copy
+def _get_activation_fn(activation):
+    if activation == "relu":
+        return F.relu
+    if activation == "gelu":
+        return F.gelu
+    if activation == "glu":
+        return F.glu
+    raise RuntimeError(F"activation should be relu/gelu, not {activation}.")
+class TransformerEncoderLayer(nn.Module):
+    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1,
+                 activation="relu"):
+        super().__init__()
+        self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout)
+        self.linear1 = nn.Linear(d_model, dim_feedforward)
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(dim_feedforward, d_model)
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+        self.activation = _get_activation_fn(activation)
+    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
+        return tensor if pos is None else tensor + pos
+    def forward(self, src,
+                    src_mask: Optional[Tensor] = None,
+                    src_key_padding_mask: Optional[Tensor] = None,
+                    pos: Optional[Tensor] = None):
+        src2 = self.norm1(src)
+        q = k = self.with_pos_embed(src2, pos)
+        src2 = self.self_attn(q, k, value=src2, attn_mask=src_mask,
+                              key_padding_mask=src_key_padding_mask)[0]
+        src = src + self.dropout1(src2)
+        src2 = self.norm2(src)
+        src2 = self.linear2(self.dropout(self.activation(self.linear1(src2))))
+        src = src + self.dropout2(src2)
+        return src
+
+
+def _get_clones(module, N):
+    return nn.ModuleList([copy.deepcopy(module) for i in range(N)])
+class TransformerEncoder(nn.Module):
+    def __init__(self, encoder_layer, num_layers, norm=None):
+        super().__init__()
+        self.layers = _get_clones(encoder_layer, num_layers)
+        self.num_layers = num_layers
+        self.norm = norm
+    def forward(self, src,
+                mask: Optional[Tensor] = None,
+                src_key_padding_mask: Optional[Tensor] = None,
+                pos: Optional[Tensor] = None):
+        output = src
+        for layer in self.layers:
+            output = layer(output, src_mask=mask,
+                           src_key_padding_mask=src_key_padding_mask, pos=pos)
+        if self.norm is not None:
+            output = self.norm(output)
+        return output
+
+
+class Transformer(nn.Module):
+    def __init__(self, d_model=512, nhead=8, num_encoder_layers=6,
+                 num_decoder_layers=6, dim_feedforward=2048, dropout=0.1,
+                 activation="relu"
+                 ):
+        super().__init__()
+        self.num_encoder_layers = num_encoder_layers
+        if num_decoder_layers > 0:
+            encoder_layer = TransformerEncoderLayer(d_model, nhead, dim_feedforward,
+                                                    dropout, activation)
+            encoder_norm = nn.LayerNorm(d_model)
+            self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm)
+
+        self._reset_parameters()
+        self.d_model = d_model
+        self.nhead = nhead
+
+    def _reset_parameters(self):
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+
+    def forward(self, src, pos_embed=None, mask=None):
+        src = src.permute(2, 0, 1)
+        if mask is not None:
+            mask = mask.flatten(1)
+        memory = self.encoder(src, src_key_padding_mask=mask, pos=pos_embed)
+        return memory.transpose(0,1)

pretrain/transformer.py

@@ -0,0 +1,172 @@
+#most of the codes below are copied from Query2label and DETR
+import torch
+import torch.nn.functional as F
+from torch import nn, Tensor
+from torch.nn import MultiheadAttention
+from typing import Optional, List
+import copy
+def _get_activation_fn(activation):
+    if activation == "relu":
+        return F.relu
+    if activation == "gelu":
+        return F.gelu
+    if activation == "glu":
+        return F.glu
+    raise RuntimeError(F"activation should be relu/gelu, not {activation}.")
+class TransformerEncoderLayer(nn.Module):
+    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1,
+                 activation="relu"):
+        super().__init__()
+        self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout)
+        self.linear1 = nn.Linear(d_model, dim_feedforward)
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(dim_feedforward, d_model)
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+        self.activation = _get_activation_fn(activation)
+    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
+        return tensor if pos is None else tensor + pos
+    def forward(self, src,
+                    src_mask: Optional[Tensor] = None,
+                    src_key_padding_mask: Optional[Tensor] = None,
+                    pos: Optional[Tensor] = None):
+        src2 = self.norm1(src)
+        q = k = self.with_pos_embed(src2, pos)
+        src2 = self.self_attn(q, k, value=src2, attn_mask=src_mask,
+                              key_padding_mask=src_key_padding_mask)[0]
+        src = src + self.dropout1(src2)
+        src2 = self.norm2(src)
+        src2 = self.linear2(self.dropout(self.activation(self.linear1(src2))))
+        src = src + self.dropout2(src2)
+        return src
+
+class TransformerDecoderLayer(nn.Module):
+    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1,
+                 activation="relu"):
+        super().__init__()
+        self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout)
+        self.multihead_attn = MultiheadAttention(d_model, nhead, dropout=dropout)
+        # Implementation of Feedforward model
+        self.linear1 = nn.Linear(d_model, dim_feedforward)
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(dim_feedforward, d_model)
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.norm3 = nn.LayerNorm(d_model)
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+        self.dropout3 = nn.Dropout(dropout)
+        self.activation = _get_activation_fn(activation)
+    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
+        return tensor if pos is None else tensor + pos
+    def forward(self, tgt, memory,
+                    tgt_mask: Optional[Tensor] = None,
+                    memory_mask: Optional[Tensor] = None,
+                    tgt_key_padding_mask: Optional[Tensor] = None,
+                    memory_key_padding_mask: Optional[Tensor] = None,
+                    pos: Optional[Tensor] = None,
+                    query_pos: Optional[Tensor] = None):
+        tgt2 = self.norm1(tgt)
+        q = k = self.with_pos_embed(tgt2, query_pos)
+        tgt2 = self.self_attn(q, k, value=tgt2, attn_mask=tgt_mask,
+                              key_padding_mask=tgt_key_padding_mask)[0]
+        tgt = tgt + self.dropout1(tgt2)
+        tgt2 = self.norm2(tgt)
+        tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt2, query_pos),
+                                   key=self.with_pos_embed(memory, pos),
+                                   value=memory, attn_mask=memory_mask,
+                                   key_padding_mask=memory_key_padding_mask)[0]
+        tgt = tgt + self.dropout2(tgt2)
+        tgt2 = self.norm3(tgt)
+        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2))))
+        tgt = tgt + self.dropout3(tgt2)
+        return tgt
+
+def _get_clones(module, N):
+    return nn.ModuleList([copy.deepcopy(module) for i in range(N)])
+class TransformerEncoder(nn.Module):
+    def __init__(self, encoder_layer, num_layers, norm=None):
+        super().__init__()
+        self.layers = _get_clones(encoder_layer, num_layers)
+        self.num_layers = num_layers
+        self.norm = norm
+    def forward(self, src,
+                mask: Optional[Tensor] = None,
+                src_key_padding_mask: Optional[Tensor] = None,
+                pos: Optional[Tensor] = None):
+        output = src
+        for layer in self.layers:
+            output = layer(output, src_mask=mask,
+                           src_key_padding_mask=src_key_padding_mask, pos=pos)
+        if self.norm is not None:
+            output = self.norm(output)
+        return output
+
+class TransformerDecoder(nn.Module):
+    def __init__(self, decoder_layer, num_layers, norm=None, return_intermediate=False):
+        super().__init__()
+        self.layers = _get_clones(decoder_layer, num_layers)
+        self.num_layers = num_layers
+        self.norm = norm
+        self.return_intermediate = return_intermediate
+    def forward(self, tgt, memory,
+                tgt_mask: Optional[Tensor] = None,
+                memory_mask: Optional[Tensor] = None,
+                tgt_key_padding_mask: Optional[Tensor] = None,
+                memory_key_padding_mask: Optional[Tensor] = None,
+                pos: Optional[Tensor] = None,
+                query_pos: Optional[Tensor] = None):
+        output = tgt
+        for layer in self.layers:
+            output = layer(output, memory, tgt_mask=tgt_mask,
+                           memory_mask=memory_mask,
+                           tgt_key_padding_mask=tgt_key_padding_mask,
+                           memory_key_padding_mask=memory_key_padding_mask,
+                           pos=pos, query_pos=query_pos)
+        if self.norm is not None:
+            output = self.norm(output)
+        return output
+
+
+class Transformer(nn.Module):
+    def __init__(self, d_model=512, nhead=8, num_encoder_layers=6,
+                 num_decoder_layers=6, dim_feedforward=2048, dropout=0.1,
+                 activation="relu"
+                 ):
+        super().__init__()
+        self.num_encoder_layers = num_encoder_layers
+        if num_decoder_layers > 0:
+            encoder_layer = TransformerEncoderLayer(d_model, nhead, dim_feedforward,
+                                                    dropout, activation)
+            encoder_norm = nn.LayerNorm(d_model)
+            self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm)
+
+        decoder_layer = TransformerDecoderLayer(d_model, nhead, dim_feedforward,
+                                                dropout, activation)
+        decoder_norm = nn.LayerNorm(d_model)
+        self.decoder = TransformerDecoder(decoder_layer, num_decoder_layers, decoder_norm)
+
+        self._reset_parameters()
+        self.d_model = d_model
+        self.nhead = nhead
+
+    def _reset_parameters(self):
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+
+    def forward(self, src, query_embed, pos_embed=None, mask=None):
+        bs, c, w = src.shape
+        src = src.permute(2, 0, 1)
+        query_embed = query_embed.transpose(0,1)
+            # .unsqueeze(1).repeat(1, bs, 1)
+        if mask is not None:
+            mask = mask.flatten(1)
+
+        memory = self.encoder(src, src_key_padding_mask=mask, pos=pos_embed)
+        hs = self.decoder(query_embed, memory, memory_key_padding_mask=mask,
+                          pos=pos_embed, query_pos=None)
+        return hs.transpose(0,1)
+        # return hs.transpose(1, 2), memory[:h * w].permute(1, 2, 0).view(bs, c, h, w)