shiji_solution2.py

import pandas as pd
import numpy as np
import warnings
import gc,os
from time import time
import datetime,random
from tqdm.notebook import tqdm
from sklearn.preprocessing import LabelEncoder
from sklearn.model_selection import KFold, StratifiedKFold,GroupKFold
from sklearn.cluster import KMeans
from sklearn.decomposition import PCA
from sklearn.feature_selection import VarianceThreshold
from sklearn.preprocessing import QuantileTransformer
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset,TensorDataset, DataLoader,RandomSampler
from iterstrat.ml_stratifiers import MultilabelStratifiedKFold
import argparse

def Parse_args():
    args = argparse.ArgumentParser()
    args.add_argument('--input_dir',
                      default='./data', help='input data path of dataset')
    args = args.parse_args()
    return args

args = Parse_args()

warnings.simplefilter('ignore')

ncompo_genes = 600
ncompo_cells = 50

def Seed_everything(seed=42):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True

Seed_everything(seed=42)

def Metric(labels,preds):
    labels = np.array(labels)
    preds = np.array(preds)
    metric = 0
    for i in range(labels.shape[1]):
        metric += (-np.mean(labels[:,i]*np.log(np.maximum(preds[:,i],1e-15))+(1-labels[:,i])*np.log(np.maximum(1-preds[:,i],1e-15))))
    return metric/labels.shape[1]

files = ['%s/test_features.csv'%args.input_dir,
         '%s/train_targets_scored.csv'%args.input_dir,
         '%s/train_features.csv'%args.input_dir,
         '%s/train_targets_nonscored.csv'%args.input_dir,
         '%s/train_drug.csv'%args.input_dir,
         '%s/sample_submission.csv'%args.input_dir]

test = pd.read_csv(files[0])
train_target = pd.read_csv(files[1])
train = pd.read_csv(files[2])
train_nonscored = pd.read_csv(files[3])
train_drug = pd.read_csv(files[4])
sub = pd.read_csv(files[5])

genes = [col for col in train.columns if col.startswith("g-")]
cells = [col for col in train.columns if col.startswith("c-")]

features = genes + cells
targets = [col for col in train_target if col!='sig_id']

ori_train = train.copy()
ctl_train = train.loc[train['cp_type']=='ctl_vehicle'].append(test.loc[test['cp_type']=='ctl_vehicle']).reset_index(drop=True)
ctl_train2 = train.loc[train['cp_type']=='ctl_vehicle'].reset_index(drop=True)

ori_test = test.copy()
ctl_test = test.loc[test['cp_type']=='ctl_vehicle'].reset_index(drop=True)

def Feature(df):
    transformers={}
    for col in tqdm(genes+cells):
        transformer = QuantileTransformer(n_quantiles=100,random_state=0, output_distribution='normal')
        transformer.fit(df[:train.shape[0]][col].values.reshape(-1,1))
        df[col] = transformer.transform(df[col].values.reshape(-1,1)).reshape(1,-1)[0]
        transformers[col]=transformer
    gene_pca = PCA(n_components = ncompo_genes,
                    random_state = 42).fit(df[genes])
    pca_genes = gene_pca.transform(df[genes])
    cell_pca = PCA(n_components = ncompo_cells,
                    random_state = 42).fit(df[cells])
    pca_cells = cell_pca.transform(df[cells])
    pca_genes = pd.DataFrame(pca_genes, columns = [f"pca_g-{i}" for i in range(ncompo_genes)])
    pca_cells = pd.DataFrame(pca_cells, columns = [f"pca_c-{i}" for i in range(ncompo_cells)])
    df = pd.concat([df, pca_genes, pca_cells], axis = 1)

    nor_var_col = [col for col in df.columns if col in ['sig_id','cp_type','cp_time','cp_dose'] or '_gt_' in col or '_lt_' in col]

    var_thresh = VarianceThreshold(0.8)
    var_cols = [col for col in df.columns if col not in ['sig_id','cp_type','cp_time','cp_dose'] and '_gt_' not in col and '_lt_' not in col]
    var_data = var_thresh.fit_transform(df[var_cols])
    df = pd.concat([df[nor_var_col],pd.DataFrame(var_data)],axis=1)
    for col in ['cp_time','cp_dose']:
        tmp = pd.get_dummies(df[col],prefix=col)
        df = pd.concat([df,tmp],axis=1)
        df.drop([col],axis=1,inplace=True)
    return df,transformers,gene_pca,cell_pca,var_thresh

tt = train.append(test).reset_index(drop=True)
tt,transformers,gene_pca,cell_pca,var_thresh = Feature(tt)
train = tt[:train.shape[0]]
test = tt[train.shape[0]:].reset_index(drop=True)

if 1:
    train_target = train_target.loc[train['cp_type']!='ctl_vehicle'].reset_index(drop=True)
    train_nonscored = train_nonscored.loc[train['cp_type']!='ctl_vehicle'].reset_index(drop=True)
    train_drug = train_drug.loc[train['cp_type']!='ctl_vehicle'].reset_index(drop=True)
    ori_train = ori_train.loc[train['cp_type']!='ctl_vehicle'].reset_index(drop=True)
    train = train.loc[train['cp_type']!='ctl_vehicle'].reset_index(drop=True)

class Model(nn.Module):
    def __init__(self, num_features, num_targets, hidden_size):
        super(Model, self).__init__()
        self.batch_norm1 = nn.BatchNorm1d(num_features)
        self.dense1 = nn.utils.weight_norm(nn.Linear(num_features, hidden_size))

        self.batch_norm2 = nn.BatchNorm1d(hidden_size)
        self.dropout2 = nn.Dropout(0.2619422201258426)
        self.dense2 = nn.utils.weight_norm(nn.Linear(hidden_size, hidden_size))

        self.batch_norm4 = nn.BatchNorm1d(hidden_size)
        self.dropout4 = nn.Dropout(0.2619422201258426)
        self.dense4 = nn.utils.weight_norm(nn.Linear(hidden_size, hidden_size))

        self.batch_norm3 = nn.BatchNorm1d(hidden_size)
        self.dropout3 = nn.Dropout(0.2619422201258426)
        self.dense3 = nn.utils.weight_norm(nn.Linear(hidden_size, num_targets))

    def forward(self, x):
        x = self.batch_norm1(x)
        x = F.leaky_relu(self.dense1(x))

        x = self.batch_norm2(x)
        x = self.dropout2(x)
        x = F.leaky_relu(self.dense2(x))

        x = self.batch_norm4(x)
        x = self.dropout4(x)
        x = F.leaky_relu(self.dense4(x))

        x = self.batch_norm3(x)
        x = self.dropout3(x)
        x = self.dense3(x)

        return x

class LabelSmoothingLoss(nn.Module):
    def __init__(self, classes, smoothing=0.0, dim=-1):
        super(LabelSmoothingLoss, self).__init__()
        self.confidence = 1.0 - smoothing
        self.smoothing = smoothing
        self.cls = classes
        self.dim = dim

    def forward(self, pred, target):
        pred = pred.log_softmax(dim=self.dim)
        with torch.no_grad():
            # true_dist = pred.data.clone()
            true_dist = torch.zeros_like(pred)
            true_dist.fill_(self.smoothing / (self.cls - 1))
            true_dist.scatter_(1, target.data.unsqueeze(1), self.confidence)
        return torch.mean(torch.sum(-true_dist * pred, dim=self.dim))

class resnetModel(nn.Module):
    def __init__(self, num_features, num_targets, hidden_size,ispretrain=False):
        super(resnetModel, self).__init__()
        self.ispretrain=ispretrain
        self.batch_norm1 = nn.BatchNorm1d(num_features)
        self.dense1 = nn.utils.weight_norm(nn.Linear(num_features, hidden_size))


        self.batch_norm2 = nn.BatchNorm1d(num_features+hidden_size)
        self.dropout2 = nn.Dropout(0.2619422201258426)
        self.dense2 = nn.utils.weight_norm(nn.Linear(num_features+hidden_size, hidden_size))
        self.batch_norm20 = nn.BatchNorm1d(hidden_size)
        self.dropout20 = nn.Dropout(0.2619422201258426)
        self.dense20 = nn.utils.weight_norm(nn.Linear(hidden_size, hidden_size))


        self.batch_norm3 = nn.BatchNorm1d(2*hidden_size)
        self.dropout3 = nn.Dropout(0.2619422201258426)
        self.dense3 = nn.utils.weight_norm(nn.Linear(2*hidden_size, hidden_size))
        self.batch_norm30 = nn.BatchNorm1d(hidden_size)
        self.dropout30 = nn.Dropout(0.2619422201258426)
        self.dense30 = nn.utils.weight_norm(nn.Linear(hidden_size, hidden_size))

        #self.batch_norm6 = nn.BatchNorm1d(2*hidden_size)
        #self.dropout6 = nn.Dropout(0.2619422201258426)
        #self.dense6 = nn.utils.weight_norm(nn.Linear(2*hidden_size, hidden_size))
        #self.batch_norm60 = nn.BatchNorm1d(hidden_size)
        #self.dropout60 = nn.Dropout(0.2619422201258426)
        #self.dense60 = nn.utils.weight_norm(nn.Linear(hidden_size, hidden_size))


        self.batch_norm4 = nn.BatchNorm1d(2*hidden_size)
        self.dropout4 = nn.Dropout(0.2619422201258426)
        if self.ispretrain:
          self.dense4 = nn.utils.weight_norm(nn.Linear(2*hidden_size, num_targets))
        else:
          self.dense5 = nn.utils.weight_norm(nn.Linear(2*hidden_size, num_targets))

    def forward(self, x):
        x1 = self.batch_norm1(x)
        x1 = F.leaky_relu(self.dense1(x1))
        x = torch.cat([x,x1],1)

        x2 = self.batch_norm2(x)
        x2 = self.dropout2(x2)
        x2 = F.leaky_relu(self.dense2(x2))
        x2 = self.batch_norm20(x2)
        x2 = self.dropout20(x2)
        x2 = F.leaky_relu(self.dense20(x2))
        x = torch.cat([x1,x2],1)

        x3 = self.batch_norm3(x)
        x3 = self.dropout3(x3)
        x3 = F.leaky_relu(self.dense3(x3))
        x3 = self.batch_norm30(x3)
        x3 = self.dropout30(x3)
        x3 = F.leaky_relu(self.dense30(x3))
        x3 = torch.cat([x2,x3],1)

        #x4 = self.batch_norm3(x)
        #x4 = self.dropout3(x4)
        #x4 = F.leaky_relu(self.dense3(x4))
        #x4 = self.batch_norm30(x4)
        #x4 = self.dropout30(x4)
        #x4 = F.leaky_relu(self.dense30(x4))
        #x4 = torch.cat([x3,x4],1)

        x4 = self.batch_norm4(x3)
        x4 = self.dropout4(x4)
        if self.ispretrain:
          x4 = self.dense4(x4)
        else:
          x4 = self.dense5(x4)
        return x4

class resnetsimpleModel(nn.Module):
    def __init__(self, num_features, num_targets, hidden_size,ispretrain=False):
        super(resnetsimpleModel, self).__init__()
        self.ispretrain=ispretrain
        self.batch_norm1 = nn.BatchNorm1d(num_features)
        self.dense1 = nn.utils.weight_norm(nn.Linear(num_features, hidden_size))


        self.batch_norm2 = nn.BatchNorm1d(num_features+hidden_size)
        self.dropout2 = nn.Dropout(0.2619422201258426)
        self.dense2 = nn.utils.weight_norm(nn.Linear(num_features+hidden_size, hidden_size))

        self.batch_norm3 = nn.BatchNorm1d(2*hidden_size)
        self.dropout3 = nn.Dropout(0.2619422201258426)
        self.dense3 = nn.utils.weight_norm(nn.Linear(2*hidden_size, hidden_size))

        self.batch_norm4 = nn.BatchNorm1d(2*hidden_size)
        self.dropout4 = nn.Dropout(0.2619422201258426)
        if self.ispretrain:
          self.dense4 = nn.utils.weight_norm(nn.Linear(2*hidden_size, num_targets))
        else:
          self.dense5 = nn.utils.weight_norm(nn.Linear(2*hidden_size, num_targets))

    def forward(self, x):
        x1 = self.batch_norm1(x)
        x1 = F.leaky_relu(self.dense1(x1))
        x = torch.cat([x,x1],1)

        x2 = self.batch_norm2(x)
        x2 = self.dropout2(x2)
        x2 = F.leaky_relu(self.dense2(x2))
        x = torch.cat([x1,x2],1)

        x3 = self.batch_norm3(x)
        x3 = self.dropout3(x3)
        x3 = self.dense3(x3)
        x3 = torch.cat([x2,x3],1)

        x3 = self.batch_norm4(x3)
        x3 = self.dropout4(x3)
        if self.ispretrain:
          x3 = self.dense4(x3)
        else:
          x3 = self.dense5(x3)
        return x3

class transModel(nn.Module):
    def __init__(self, num_features, num_targets, hidden_size):
        super(transModel, self).__init__()
        d_model = 20
        self.batch_norm1 = nn.BatchNorm1d(num_features)
        self.dropout1 = nn.Dropout(0.2619422201258426)
        self.dense1 = nn.utils.weight_norm(nn.Linear(num_features, hidden_size))

        encoder_layer = nn.TransformerEncoderLayer(d_model=d_model, nhead=4, dropout=0.75)
        # self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=3, norm=nn.BatchNorm1d(hidden_size // d_model, eps=1e-5))
        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=2, norm=nn.LayerNorm(d_model, eps=1e-5))
        #self.transformer_encoder2 = nn.TransformerEncoder(encoder_layer, num_layers=2, norm=nn.LayerNorm(d_model, eps=1e-5))

        self.batch_norm2 = nn.BatchNorm1d(hidden_size)
        self.dropout2 = nn.Dropout(0.2619422201258426)
        self.dense2 = nn.utils.weight_norm(nn.Linear(hidden_size, hidden_size))

        self.batch_norm3 = nn.BatchNorm1d(hidden_size)
        self.dropout3 = nn.Dropout(0.2619422201258426)
        self.dense3 = nn.utils.weight_norm(nn.Linear(hidden_size, num_targets))

    def forward(self, x):
        x = self.batch_norm1(x)
        x = self.dropout1(x)
        x = F.leaky_relu(self.dense1(x))
        x = x.view(x.shape[0], -1, 20)
        x = F.leaky_relu(self.transformer_encoder(x))
        x = x.view(x.shape[0], -1)

        x = self.batch_norm2(x)
        x = self.dropout2(x)
        x = F.leaky_relu(self.dense2(x))

        x = self.batch_norm3(x)
        x = self.dropout3(x)
        x = self.dense3(x)

        return x

import torch
from torch.nn.modules.loss import _WeightedLoss
import torch.nn.functional as F

class SmoothBCEwLogits(_WeightedLoss):
    def __init__(self, weight=None, reduction='mean', smoothing=0.0):
        super().__init__(weight=weight, reduction=reduction)
        self.smoothing = smoothing
        self.weight = weight
        self.reduction = reduction

    @staticmethod
    def _smooth(targets:torch.Tensor, n_labels:int, smoothing=0.0):
        assert 0 <= smoothing < 1
        with torch.no_grad():
            targets = targets * (1.0 - smoothing) + 0.5 * smoothing
        return targets

    def forward(self, inputs, targets):
        targets = SmoothBCEwLogits._smooth(targets, inputs.size(-1),
            self.smoothing)
        loss = F.binary_cross_entropy_with_logits(inputs, targets,self.weight)

        if  self.reduction == 'sum':
            loss = loss.sum()
        elif  self.reduction == 'mean':
            loss = loss.mean()

        return loss

def Ctl_augment(train,target,include_test=0):
    if include_test==0:
        ctl_aug=ctl_train2.copy()
    if include_test==1:
        ctl_aug=ctl_train.copy()
    aug_trains = []
    aug_targets = []
    for t in [24,48,72]:
        for d in ['D1','D2']:
            for _ in range(2):
                train1 = train.loc[(train['cp_time']==t)&(train['cp_dose']==d)]
                target1 = target.loc[(train['cp_time']==t)&(train['cp_dose']==d)]
                ctl1 = ctl_aug.loc[(ctl_aug['cp_time']==t)&(ctl_aug['cp_dose']==d)].sample(train1.shape[0],replace=True)
                ctl2 = ctl_aug.loc[(ctl_aug['cp_time']==t)&(ctl_aug['cp_dose']==d)].sample(train1.shape[0],replace=True)
                train1[genes+cells] = train1[genes+cells].values + ctl1[genes+cells].values - ctl2[genes+cells].values
                aug_train = train1.merge(target1,how='left',on='sig_id')
                aug_trains.append(aug_train[['cp_time','cp_dose']+genes+cells])
                aug_targets.append(aug_train[targets])
    df = pd.concat(aug_trains).reset_index(drop=True)
    target = pd.concat(aug_targets).reset_index(drop=True)
    for col in tqdm(genes+cells):
        df[col] = transformers[col].transform(df[col].values.reshape(-1,1)).reshape(1,-1)[0]
    pca_genes = gene_pca.transform(df[genes])
    pca_cells = cell_pca.transform(df[cells])
    pca_genes = pd.DataFrame(pca_genes, columns = [f"pca_g-{i}" for i in range(ncompo_genes)])
    pca_cells = pd.DataFrame(pca_cells, columns = [f"pca_c-{i}" for i in range(ncompo_cells)])
    df = pd.concat([df, pca_genes, pca_cells], axis = 1)
    for col in ['cp_time','cp_dose']:
        tmp = pd.get_dummies(df[col],prefix=col)
        df = pd.concat([df,tmp],axis=1)
        df.drop([col],axis=1,inplace=True)
    xs = df[train_cols].values
    ys = target[targets]
    #ys_ns = target[targets_ns]
    return xs,ys#,ys_ns

def Ctl_augment2(train,target,include_test=0):
    if include_test==0:
        ctl_aug=ctl_train2.copy()
    if include_test==1:
        ctl_aug=ctl_train.copy()
    aug_trains = []
    aug_targets = []
    for t in [24,48,72]:
        for d in ['D1','D2']:
            for _ in range(1):
                train1 = train.loc[(train['cp_time']==t)&(train['cp_dose']==d)]
                target1 = target.loc[(train['cp_time']==t)&(train['cp_dose']==d)]
                ctl1 = ctl_aug.loc[(ctl_aug['cp_time']==t)&(ctl_aug['cp_dose']==d)].sample(train1.shape[0],replace=True)
                ctl2 = ctl_aug.loc[(ctl_aug['cp_time']==t)&(ctl_aug['cp_dose']==d)].sample(train1.shape[0],replace=True)
                ctl3 = ctl_aug.loc[(ctl_aug['cp_time']==t)&(ctl_aug['cp_dose']==d)].sample(train1.shape[0],replace=True)
                ctl4 = ctl_aug.loc[(ctl_aug['cp_time']==t)&(ctl_aug['cp_dose']==d)].sample(train1.shape[0],replace=True)
                train1[genes+cells] = train1[genes+cells].values + ctl1[genes+cells].values + ctl2[genes+cells].values -ctl3[genes+cells].values - ctl4[genes+cells].values
                aug_train = train1.merge(target1,how='left',on='sig_id')
                aug_trains.append(aug_train[['cp_time','cp_dose']+genes+cells])
                aug_targets.append(aug_train[targets])
    df = pd.concat(aug_trains).reset_index(drop=True)
    target = pd.concat(aug_targets).reset_index(drop=True)
    for col in tqdm(genes+cells):
        df[col] = transformers[col].transform(df[col].values.reshape(-1,1)).reshape(1,-1)[0]
    pca_genes = gene_pca.transform(df[genes])
    pca_cells = cell_pca.transform(df[cells])
    pca_genes = pd.DataFrame(pca_genes, columns = [f"pca_g-{i}" for i in range(ncompo_genes)])
    pca_cells = pd.DataFrame(pca_cells, columns = [f"pca_c-{i}" for i in range(ncompo_cells)])
    df = pd.concat([df, pca_genes, pca_cells], axis = 1)
    for col in ['cp_time','cp_dose']:
        tmp = pd.get_dummies(df[col],prefix=col)
        df = pd.concat([df,tmp],axis=1)
        df.drop([col],axis=1,inplace=True)
    xs = df[train_cols].values
    ys = target[targets]
    #ys_ns = target[targets_ns]
    return xs,ys#,ys_ns

def Ctl_augment_new(train,target,include_test=0):
    if include_test==0:
        ctl_aug=ctl_train2.copy()
    if include_test==1:
        ctl_aug=ctl_train.copy()
    aug_trains = []
    aug_targets = []
    for _ in range(3):
          train1 = train.copy()
          target1 = target.copy()
          ctl1 = ctl_train.sample(train1.shape[0],replace=True).reset_index(drop=True)#.loc[(ctl_train['cp_time']==t)&(ctl_train['cp_dose']==d)]
          ctl2 = ctl_train.sample(train1.shape[0],replace=True).reset_index(drop=True)

          ctl3 = ctl_train.sample(train1.shape[0],replace=True).reset_index(drop=True)#.loc[(ctl_train['cp_time']==t)&(ctl_train['cp_dose']==d)]
          ctl4 = ctl_train.sample(train1.shape[0],replace=True).reset_index(drop=True)
          mask_index1 = list(np.random.choice(ctl3.index.tolist(),int(ctl3.shape[0]*0.4),replace=False))
          ctl3.loc[mask_index1,genes+cells] = 0.0
          ctl4.loc[mask_index1,genes+cells] = 0.0

          ctl5 = ctl_train.sample(train1.shape[0],replace=True).reset_index(drop=True)#.loc[(ctl_train['cp_time']==t)&(ctl_train['cp_dose']==d)]
          ctl6 = ctl_train.sample(train1.shape[0],replace=True).reset_index(drop=True)
          mask_index2 = list(np.random.choice(list(set(ctl5.index)-set(mask_index1)),int(ctl5.shape[0]*0.3),replace=False))
          ctl5.loc[mask_index1+mask_index2,genes+cells] = 0.0
          ctl6.loc[mask_index1+mask_index2,genes+cells] = 0.0

          train1[genes+cells] = train1[genes+cells].values + ctl1[genes+cells].values - ctl2[genes+cells].values \
                              + ctl3[genes+cells].values - ctl4[genes+cells].values + ctl5[genes+cells].values - ctl6[genes+cells].values

          aug_train = train1.merge(target1,how='left',on='sig_id')
          aug_trains.append(aug_train[['cp_time','cp_dose']+genes+cells])
          aug_targets.append(aug_train[targets])

    df = pd.concat(aug_trains).reset_index(drop=True)
    target = pd.concat(aug_targets).reset_index(drop=True)
    for col in tqdm(genes+cells):
        df[col] = transformers[col].transform(df[col].values.reshape(-1,1)).reshape(1,-1)[0]
    pca_genes = gene_pca.transform(df[genes])
    pca_cells = cell_pca.transform(df[cells])
    pca_genes = pd.DataFrame(pca_genes, columns = [f"pca_g-{i}" for i in range(ncompo_genes)])
    pca_cells = pd.DataFrame(pca_cells, columns = [f"pca_c-{i}" for i in range(ncompo_cells)])
    df = pd.concat([df, pca_genes, pca_cells], axis = 1)

    nor_var_col = [col for col in df.columns if col in ['sig_id','cp_type','cp_time','cp_dose'] or '_gt_' in col or '_lt_' in col]
    var_cols = [col for col in df.columns if col not in ['sig_id','cp_type','cp_time','cp_dose'] and '_gt_' not in col and '_lt_' not in col]
    var_data = var_thresh.transform(df[var_cols])
    df = pd.concat([df[nor_var_col],pd.DataFrame(var_data)],axis=1)

    for col in ['cp_time','cp_dose']:
        tmp = pd.get_dummies(df[col],prefix=col)
        df = pd.concat([df,tmp],axis=1)
        df.drop([col],axis=1,inplace=True)
    xs = df[train_cols].values
    ys = target[targets]
    #ys_ns = target[targets_ns]
    return xs,ys#,ys_ns

class MoADataset:
    def __init__(self, features, targets,noise=0.1,val=0):
        self.features = features
        self.targets = targets
        self.noise = noise
        self.val = val

    def __len__(self):
        return (self.features.shape[0])

    def __getitem__(self, idx):
        sample = self.features[idx, :].copy()

        if 0 and np.random.rand()<0.3 and not self.val:
            sample = self.swap_sample(sample)

        dct = {
            'x' : torch.tensor(sample, dtype=torch.float),
            'y' : torch.tensor(self.targets[idx, :], dtype=torch.float)
        }
        return dct

    def swap_sample(self,sample):
            #print(sample.shape)
            num_samples = self.features.shape[0]
            num_features = self.features.shape[1]
            if len(sample.shape) == 2:
                batch_size = sample.shape[0]
                random_row = np.random.randint(0, num_samples, size=batch_size)
                for i in range(batch_size):
                    random_col = np.random.rand(num_features) < self.noise
                    #print(random_col)
                    sample[i, random_col] = self.features[random_row[i], random_col]
            else:
                batch_size = 1

                random_row = np.random.randint(0, num_samples, size=batch_size)


                random_col = np.random.rand(num_features) < self.noise
                #print(random_col)
                #print(random_col)

                sample[ random_col] = self.features[random_row, random_col]

            return sample

class TestDataset:
    def __init__(self, features):
        self.features = features

    def __len__(self):
        return (self.features.shape[0])

    def __getitem__(self, idx):
        dct = {
            'x' : torch.tensor(self.features[idx, :], dtype=torch.float)
        }
        return dct

device = ('cuda' if torch.cuda.is_available() else 'cpu')
EPOCHS1 = 29
EPOCHS = 23
trn_loss_=[]
def train_and_predict(features, sub, aug,  folds=5, seed=817119,lr=1/90.0/3.5*3,weight_decay=1e-5/3):
    oof = train[['sig_id']]
    for t in targets:
        oof[t] = 0.0
    preds = []
    test_X = test[features].values
    test_data_loader = DataLoader(dataset=TensorDataset(torch.Tensor(test_X)),batch_size=1024,shuffle=False)
    eval_train_loss = 0.0
    for fold, (trn_ind, val_ind) in enumerate(MultilabelStratifiedKFold(n_splits = folds, shuffle=True, random_state=seed)\
                                              .split(train, train_target[targets])):
        train_X = train.loc[trn_ind,features].values
        train_Y = train_target.loc[trn_ind,targets].values
        eval_train_Y = train_target.loc[trn_ind,targets].values
        eval_train_dataset = MoADataset(train_X, eval_train_Y)
        eval_train_data_loader = torch.utils.data.DataLoader(eval_train_dataset, batch_size=128, shuffle=False)

        valid_X = train.loc[val_ind,features].values
        valid_Y = train_target.loc[val_ind,targets].values
        valid_dataset = MoADataset(valid_X, valid_Y,val=1)
        valid_data_loader = torch.utils.data.DataLoader(valid_dataset, batch_size=1024, shuffle=False)

        aug_X,aug_Y = Ctl_augment_new(ori_train.loc[trn_ind],train_target.loc[trn_ind],include_test=1)
        train_X_ = np.concatenate([train_X,aug_X],axis=0)
        train_Y_ = np.concatenate([train_Y,aug_Y],axis=0)

        train_dataset = MoADataset(train_X_, train_Y_)
        train_data_loader = torch.utils.data.DataLoader(train_dataset, batch_size=128, shuffle=True)

        model = resnetModel(len(features),len(targets),1500)
        model.to(device)


        optimizer = torch.optim.Adam(model.parameters(),betas=(0.9, 0.99), lr=1e-3, weight_decay=weight_decay,eps=1e-5)
        scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer=optimizer, pct_start=0.1, div_factor=1e3,
                                              max_lr=lr, epochs=EPOCHS1, steps_per_epoch=len(train_data_loader))

        loss_fn = nn.BCEWithLogitsLoss()
        loss_tr = SmoothBCEwLogits(smoothing =0.001)

        best_valid_metric = 1e9
        not_improve_epochs = 0
        for epoch in range(EPOCHS1):
            # train
            train_loss = 0.0
            train_num = 0
            for data in (train_data_loader):
                optimizer.zero_grad()
                x,y = data['x'].to(device),data['y'].to(device)
                outputs = model(x)
                loss = loss_tr(outputs, y)
                loss.backward()
                optimizer.step()
                scheduler.step()
                train_num += x.shape[0]
                train_loss += (loss.item()*x.shape[0])

            train_loss /= train_num
            # eval
            model.eval()
            valid_loss = 0.0
            valid_num = 0
            for data in (valid_data_loader):
                x,y = data['x'].to(device),data['y'].to(device)
                outputs = model(x)
                loss = loss_fn(outputs, y)
                valid_num += x.shape[0]
                valid_loss += (loss.item()*x.shape[0])
            valid_loss /= valid_num
            t_preds = []
            for data in (test_data_loader):
                x = data[0].to(device)
                with torch.no_grad():
                    outputs = model(x)
                t_preds.extend(list(outputs.sigmoid().cpu().detach().numpy()))
            pred_mean = np.mean(t_preds)
            if valid_loss < best_valid_metric:
                torch.save(model.state_dict(),'./model/model_resnet2_fold%s'%fold+'_'+str(seed)+'.ckpt')
                not_improve_epochs = 0
                best_valid_metric = valid_loss
                print('[epoch %s] lr: %.6f, train_loss: %.6f, valid_metric: %.6f, pred_mean:%.6f'%(epoch,optimizer.param_groups[0]['lr'],train_loss,valid_loss,pred_mean))
                trn_loss_.append(train_loss)
            else:
                not_improve_epochs += 1
                print('[epoch %s] lr: %.6f, train_loss: %.6f, valid_metric: %.6f, pred_mean:%.6f, NIE +1 ---> %s'%(epoch,optimizer.param_groups[0]['lr'],train_loss,valid_loss,pred_mean,not_improve_epochs))
                if not_improve_epochs >= 30 and epoch>15:
                    break
            model.train()
            if epoch!=28:
                aug_X,aug_Y = Ctl_augment_new(ori_train.loc[trn_ind],train_target.loc[trn_ind],include_test=1)
                train_X_ = np.concatenate([train_X,aug_X],axis=0)
                train_Y_ = np.concatenate([train_Y,aug_Y],axis=0)
                train_dataset = MoADataset(train_X_, train_Y_)
                train_data_loader = torch.utils.data.DataLoader(train_dataset, batch_size=128, shuffle=True)

        state_dict = torch.load('./model/model_resnet2_fold%s'%fold+'_'+str(seed)+'.ckpt', torch.device("cuda" if torch.cuda.is_available() else "cpu") )
        model.load_state_dict(state_dict)
        model.eval()

        valid_preds = []
        for data in tqdm(valid_data_loader):
            x,y = data['x'].to(device),data['y'].to(device)
            with torch.no_grad():
                outputs = model(x)
            valid_preds.extend(list(outputs.cpu().detach().numpy()))
        oof.loc[val_ind,targets] = 1 / (1+np.exp(-np.array(valid_preds)))
        t_preds = []
        for data in tqdm(test_data_loader):
            x = data[0].to(device)
            with torch.no_grad():
                outputs = model(x)
            t_preds.extend(list(outputs.sigmoid().cpu().detach().numpy()))
        print(np.mean(t_preds))
        preds.append(t_preds)
        train_preds=[]

        for data in (eval_train_data_loader):
            x =  data['x'].to(device)
            with torch.no_grad():
                outputs = model(x)
            train_preds.extend(list(outputs.sigmoid().cpu().detach().numpy()))
        train_loss = Metric(eval_train_Y,train_preds)
        eval_train_loss += train_loss
        print('eval_train_loss:',train_loss)

    sub[targets] = np.array(preds).mean(axis=0)
    return oof,sub



train_cols = [col for col in train.columns if col not in ['sig_id','cp_type']]

Seed_everything(0)
oof,sub = train_and_predict(train_cols,sub.copy(),aug=True,seed=0,lr=1/90.0/2,weight_decay=1e-5/2.7)

outputs = []
for seed in [1,2,3]:
    Seed_everything(seed)
    outputs.append(train_and_predict(train_cols,sub.copy(),aug=True,seed=seed,lr=1/90.0/2,weight_decay=1e-5/2.7))

for output in outputs:
    oof[targets] += output[0][targets]
    sub[targets] += output[1][targets]
oof[targets] /= (1+len(outputs))
sub[targets] /= (1+len(outputs))

valid_metric = Metric(train_target[targets].values,oof[targets].values)
print('oof mean:%.6f,sub mean:%.6f,valid metric:%.6f'%(oof[targets].mean().mean(),sub[targets].mean().mean(),valid_metric))
sub.loc[test['cp_type']=='ctl_vehicle',targets] = 0.0
sub.to_csv('./shiji_submission2.csv',index=False)