classification_pytorch_dummy.py

# -*- coding: utf-8 -*-
"""classification pytorch.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/137kv2KGkMfqJf70XUlKDDlYraF_R26Ve
"""

from google.colab import files
files.upload()
#kaggle json upload

!pip install -q kaggle
!pip install -U git+https://github.com/albu/albumentations1
!mkdir -p ~/.kaggle
!cp kaggle.json ~/.kaggle/
!kaggle competitions download -c severstal-steel-defect-detection
!unzip train_images.zip -d train_images

import numpy as np # linear algebra
import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)
import warnings 
warnings.filterwarnings('ignore')

# Input data files are available in the "../input/" directory.
# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory

import os
import cv2
import os.path as osp
import time
import torch
import torch.nn as nn
import torch.optim as optim 
from torch.optim import lr_scheduler
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, models
from sklearn.model_selection import train_test_split
from tqdm import tqdm_notebook as tqdm
from PIL import Image
from torch.optim.lr_scheduler import ReduceLROnPlateau

import matplotlib.pyplot as plt

# now we train a binary classifier to check whether one image is defective or not 
TRAINVAL_ANNOT = "train.csv.zip"
TRAINVAL_IMAGE_ROOT = "train_images"

from albumentations import (HorizontalFlip, ShiftScaleRotate, Normalize, CropNonEmptyMaskIfExists, Resize, Compose,RandomCrop,
                            RandomBrightnessContrast, VerticalFlip, RandomBrightness, RandomContrast)
from albumentations.pytorch import ToTensor

# Commented out IPython magic to ensure Python compatibility.
from sklearn.model_selection import StratifiedKFold, KFold

def get_annot(annot_path):
    trainval_annot = pd.read_csv(TRAINVAL_ANNOT)
    trainval_annot['ImageId'] = trainval_annot['ImageId_ClassId'].apply(lambda x: x.split("_")[0])
    trainval_annot['ClassId'] = trainval_annot['ImageId_ClassId'].apply(lambda x: x.split("_")[1])
    trainval_annot['HasMask'] = trainval_annot['EncodedPixels'].notnull().astype('int')
    trainval_annot = trainval_annot[['ImageId', 'ClassId', 'HasMask']].groupby(['ImageId', 'ClassId']).max().unstack().sort_index(axis=1).reset_index()
    trainval_annot.columns = ['ImageId', '1', '2', '3', '4']
    kfold = KFold(10, shuffle=True,
                          random_state=69)  # StratifiedKFold(total_folds, shuffle=True, random_state=69)
    train_idx, val_idx = list(kfold.split(trainval_annot))[0]  # , df["defects"]
    train_annot, val_annot = trainval_annot.iloc[train_idx], trainval_annot.iloc[val_idx]
    #train_annot, val_annot = train_test_split(trainval_annot, test_size=0.10)
    print("{}/{} images for train/val.".format(len(train_annot), len(val_annot)))
    return {"train": train_annot, "val": val_annot}

def prev_get_transform(phase):
    list_transforms = []
    if phase == "train":
          list_transforms.extend(
              [RandomCrop(256, 1024, p=1),
                HorizontalFlip(p=0.5),
                VerticalFlip(p=0.5),
                RandomBrightnessContrast(p=0.1, brightness_limit=0.1, contrast_limit=0.1)
                ])

    list_transforms.extend(
        [
            Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], p=1),
            ToTensor()
        ]
    )
    list_trfms = Compose(list_transforms)
    return list_trfms

def get_transform(phase):
    list_transform = []
    if phase == 'train':
        list_transform.extend([
            transforms.Resize((256, 256+128)),
            transforms.RandomHorizontalFlip(p=0.5),
            transforms.RandomVerticalFlip(p=0.5),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])
    else:
        list_transform.extend([
            transforms.Resize((256, 256+128)),
            # transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])
    return transforms.Compose(list_transform)

def criterion(logit, truth, weight=None):
    batch_size,num_class = logit.shape
    assert(logit.shape==truth.shape)

    loss = F.binary_cross_entropy_with_logits(logit, truth, reduction='none')

    if weight is None:
        loss = loss.mean()

    else:
        pos = (truth>0.5).float()
        neg = (truth<0.5).float()
        pos_sum = pos.sum().item() + 1e-12
        neg_sum = neg.sum().item() + 1e-12
        loss = (weight[1]*pos*loss/pos_sum + weight[0]*neg*loss/neg_sum).sum()
        #raise NotImplementedError

    return loss

class SteelDataset(Dataset):
    def __init__(self, annot, image_folder, phase):
        self.annot = annot
        self.image_folder = image_folder
        self.phase = phase
        self.transform = get_transform(phase)
        
    def __getitem__(self, index):
        row = self.annot.iloc[index, :]
        image_path = osp.join(self.image_folder, row['ImageId'])
        target = torch.tensor([row['1'], row['2'], row['3'], row['4']], dtype=float)
        image = Image.open(image_path)
        image = self.transform(image)
        return image, target
    
    def __len__(self):
        return len(self.annot)
    
    
def get_dataloader(annot, image_folder, phase, batch_size=16, num_workers=4):
    dataset = SteelDataset(annot, image_folder, phase)
    return DataLoader(dataset, batch_size=batch_size, num_workers=num_workers)

def get_model(model_name):
    model = ResidualNet('ImageNet', 50, 4, 'cbam')
    # model = models.__dict__[model_name]()
    # if model_name.startswith("resnet"):
    #     in_features = model.fc.in_features
    #     model.fc = nn.Linear(in_features, 4)
    # else:
    #     raise KeyError("Only support resnet!")
    return model

class Trainer:
    def __init__(self,  model_name="resnet34", pretrained=False, epochs=1):
        self.lr = 5e-4
        self.threshold = 0.5
        self.best_acc = 0.0
        self.device = torch.device("cuda:0")
        self.model_path = "/content/drive/My Drive/weights/model_resne50.pth"
        self.pretrained = pretrained
        self.pretrained_model_path = osp.join("../input/severstal-binary-classifier/", self.model_path)
        self.num_epochs = epochs
        self.annot_path = TRAINVAL_ANNOT
        self.image_folder = TRAINVAL_IMAGE_ROOT
        self.phases = ['train', 'val']
        self.batch_sizes = {'train': 16, 'val': 16}#
        self.model = get_model(model_name).to(self.device)
        # self.optimizer = optim.SGD(self.model.parameters(), lr=self.lr, momentum=0.9, weight_decay=5e-4)
        self.optimizer = optim.Adam(self.model.parameters(), lr=self.lr)
        self.scheduler = lr_scheduler.StepLR(self.optimizer, step_size=5, gamma=0.5) #ReduceLROnPlateau(self.optimizer, factor=0.9, mode="min", patience=3, verbose=True)
        self.criterion = criterion#nn.CrossEntropyLoss()
        self.annots = get_annot(self.annot_path)
        self.dataloaders = {phase: get_dataloader(self.annots[phase], self.image_folder, 
                                phase, self.batch_sizes[phase]) for phase in self.phases}
        self.losses = {phase: [] for phase in self.phases}
        self.accuracies = {phase: [] for phase in self.phases}
        
    def forward(self, inputs, targets):   
        inputs = inputs.to(self.device, dtype=torch.float)
        targets = targets.to(self.device, dtype=torch.float)
        outputs = self.model(inputs)
        #targets = targets.unsqueeze(1).float()
        loss = self.criterion(outputs, targets)
        return loss, outputs
        
    def iterate(self, epoch, phase):
        start = time.time()
        print("Epoch: {} | Phase: {}".format(epoch, phase))
        self.model.train(phase == 'train')
        dataloader = self.dataloaders[phase]
        running_loss = 0.0
        running_corrects = 0
        self.optimizer.zero_grad()
        tk = tqdm(dataloader, total=len(dataloader))
        torch.set_grad_enabled(phase == 'train')
        for idx, batch in enumerate(tk):
            inputs, targets = batch
            loss, outputs = self.forward(inputs, targets)
            if phase == 'train':
                loss.backward()
                self.optimizer.step()
                self.optimizer.zero_grad()
            preds = (torch.sigmoid(outputs) > self.threshold).reshape_as(targets).int().cpu()
            running_corrects += (preds == targets.int()).sum().item()
            running_loss += loss.item() * inputs.size(0)
            tk.set_postfix(loss=running_loss / ((idx+1) * self.batch_sizes[phase]), running_corrects=running_corrects/(4*self.batch_sizes[phase]*(idx+1)))
            tk.update()
        torch.set_grad_enabled(phase == 'train')
        
        dataset_size = len(dataloader.dataset)
        running_loss /= dataset_size
        running_acc = running_corrects/ (dataset_size*4)
        self.losses[phase].append(running_loss)
        self.accuracies[phase].append(running_acc)
        end  = time.time()
        time_elapsed = int(end - start)
        print("Finished in {} mins and {} secs, loss: {:.3f}, acc: {:.3f}%".format(
            time_elapsed // 60, time_elapsed % 60, running_loss, running_acc * 100))
        return running_acc
        
    def save_model(self, epoch):
        state = {
            "epoch": epoch,
            "best_acc": self.best_acc, 
            "state_dict": self.model.state_dict(),
            "optimizer_state_dict": self.optimizer.state_dict()
        }
        print("****** Find new optimal model, saving to disk ******")
        torch.save(state, self.model_path)
        return 
    
    def summary(self):
        print("Training finished. Best val acc: {:.3f}%".format(self.best_acc))
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 8))
        len_x = len(self.losses['train'])
        ax1.plot(range(len_x), self.losses['train'], label="train loss")
        ax1.plot(range(len_x), self.losses['val'], label="val loss")
        #ax1.title("loss curve");   ax1.xlabel("epoch");   ax1.ylabel("loss")
        ax2.plot(range(len_x), self.accuracies['train'], label='train acc')
        ax2.plot(range(len_x), self.accuracies['val'], label='val acc')
        #ax2.title("accuracy curve");  ax2.xlabel("epoch");  ax2.ylabel("accuracy")
        plt.show()
    
    def start(self):
        resume_epoch = 0
        if self.pretrained or osp.exists(self.model_path):
            state = torch.load(self.pretrained_model_path) if self.pretrained else torch.load(self.model_path)
            resume_epoch = state['epoch'] + 1
            self.best_acc = state['best_acc']
            self.optimizer.load_state_dict(state['optimizer_state_dict'])
            self.model.load_state_dict(state['state_dict'])
            print("Load checkpoint from {}, resume training from {} epoch with best acc {}%".format(
                    self.model_path, resume_epoch, self.best_acc * 100))
            
        for epoch in range(resume_epoch, self.num_epochs):
            train_acc = self.iterate(epoch, 'train')
            self.scheduler.step()
            val_acc = self.iterate(epoch, 'val')
            if val_acc > self.best_acc:
                self.best_acc = val_acc
                self.save_model(epoch)
        self.summary()
# %matplotlib inline
trainer = Trainer("resnet50", pretrained=False, epochs = 100)
trainer.start()

import os
os.listdir('/content/drive/My Drive')

import torch
import math
import torch.nn as nn
import torch.nn.functional as F

class Flatten(nn.Module):
    def forward(self, x):
        return x.view(x.size(0), -1)
class ChannelGate(nn.Module):
    def __init__(self, gate_channel, reduction_ratio=16, num_layers=1):
        super(ChannelGate, self).__init__()
        self.gate_activation = gate_activation
        self.gate_c = nn.Sequential()
        self.gate_c.add_module( 'flatten', Flatten() )
        gate_channels = [gate_channel]
        gate_channels += [gate_channel // reduction_ratio] * num_layers
        gate_channels += [gate_channel]
        for i in range( len(gate_channels) - 2 ):
            self.gate_c.add_module( 'gate_c_fc_%d'%i, nn.Linear(gate_channels[i], gate_channels[i+1]) )
            self.gate_c.add_module( 'gate_c_bn_%d'%(i+1), nn.BatchNorm1d(gate_channels[i+1]) )
            self.gate_c.add_module( 'gate_c_relu_%d'%(i+1), nn.ReLU() )
        self.gate_c.add_module( 'gate_c_fc_final', nn.Linear(gate_channels[-2], gate_channels[-1]) )
    def forward(self, in_tensor):
        avg_pool = F.avg_pool2d( in_tensor, in_tensor.size(2), stride=in_tensor.size(2) )
        return self.gate_c( avg_pool ).unsqueeze(2).unsqueeze(3).expand_as(in_tensor)

class SpatialGate(nn.Module):
    def __init__(self, gate_channel, reduction_ratio=16, dilation_conv_num=2, dilation_val=4):
        super(SpatialGate, self).__init__()
        self.gate_s = nn.Sequential()
        self.gate_s.add_module( 'gate_s_conv_reduce0', nn.Conv2d(gate_channel, gate_channel//reduction_ratio, kernel_size=1))
        self.gate_s.add_module( 'gate_s_bn_reduce0',	nn.BatchNorm2d(gate_channel//reduction_ratio) )
        self.gate_s.add_module( 'gate_s_relu_reduce0',nn.ReLU() )
        for i in range( dilation_conv_num ):
            self.gate_s.add_module( 'gate_s_conv_di_%d'%i, nn.Conv2d(gate_channel//reduction_ratio, gate_channel//reduction_ratio, kernel_size=3, \
						padding=dilation_val, dilation=dilation_val) )
            self.gate_s.add_module( 'gate_s_bn_di_%d'%i, nn.BatchNorm2d(gate_channel//reduction_ratio) )
            self.gate_s.add_module( 'gate_s_relu_di_%d'%i, nn.ReLU() )
        self.gate_s.add_module( 'gate_s_conv_final', nn.Conv2d(gate_channel//reduction_ratio, 1, kernel_size=1) )
    def forward(self, in_tensor):
        return self.gate_s( in_tensor ).expand_as(in_tensor)
class BAM(nn.Module):
    def __init__(self, gate_channel):
        super(BAM, self).__init__()
        self.channel_att = ChannelGate(gate_channel)
        self.spatial_att = SpatialGate(gate_channel)
    def forward(self,in_tensor):
        att = 1 + F.sigmoid( self.channel_att(in_tensor) * self.spatial_att(in_tensor) )
        return att * in_tensor
import torch
import math
import torch.nn as nn
import torch.nn.functional as F

class BasicConv(nn.Module):
    def __init__(self, in_planes, out_planes, kernel_size, stride=1, padding=0, dilation=1, groups=1, relu=True, bn=True, bias=False):
        super(BasicConv, self).__init__()
        self.out_channels = out_planes
        self.conv = nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)
        self.bn = nn.BatchNorm2d(out_planes,eps=1e-5, momentum=0.01, affine=True) if bn else None
        self.relu = nn.ReLU() if relu else None

    def forward(self, x):
        x = self.conv(x)
        if self.bn is not None:
            x = self.bn(x)
        if self.relu is not None:
            x = self.relu(x)
        return x

class Flatten(nn.Module):
    def forward(self, x):
        return x.view(x.size(0), -1)

class ChannelGate(nn.Module):
    def __init__(self, gate_channels, reduction_ratio=16, pool_types=['avg', 'max']):
        super(ChannelGate, self).__init__()
        self.gate_channels = gate_channels
        self.mlp = nn.Sequential(
            Flatten(),
            nn.Linear(gate_channels, gate_channels // reduction_ratio),
            nn.ReLU(),
            nn.Linear(gate_channels // reduction_ratio, gate_channels)
            )
        self.pool_types = pool_types
    def forward(self, x):
        channel_att_sum = None
        for pool_type in self.pool_types:
            if pool_type=='avg':
                avg_pool = F.avg_pool2d( x, (x.size(2), x.size(3)), stride=(x.size(2), x.size(3)))
                channel_att_raw = self.mlp( avg_pool )
            elif pool_type=='max':
                max_pool = F.max_pool2d( x, (x.size(2), x.size(3)), stride=(x.size(2), x.size(3)))
                channel_att_raw = self.mlp( max_pool )
            elif pool_type=='lp':
                lp_pool = F.lp_pool2d( x, 2, (x.size(2), x.size(3)), stride=(x.size(2), x.size(3)))
                channel_att_raw = self.mlp( lp_pool )
            elif pool_type=='lse':
                # LSE pool only
                lse_pool = logsumexp_2d(x)
                channel_att_raw = self.mlp( lse_pool )

            if channel_att_sum is None:
                channel_att_sum = channel_att_raw
            else:
                channel_att_sum = channel_att_sum + channel_att_raw

        scale = F.sigmoid( channel_att_sum ).unsqueeze(2).unsqueeze(3).expand_as(x)
        return x * scale

def logsumexp_2d(tensor):
    tensor_flatten = tensor.view(tensor.size(0), tensor.size(1), -1)
    s, _ = torch.max(tensor_flatten, dim=2, keepdim=True)
    outputs = s + (tensor_flatten - s).exp().sum(dim=2, keepdim=True).log()
    return outputs

class ChannelPool(nn.Module):
    def forward(self, x):
        return torch.cat( (torch.max(x,1)[0].unsqueeze(1), torch.mean(x,1).unsqueeze(1)), dim=1 )

class SpatialGate(nn.Module):
    def __init__(self):
        super(SpatialGate, self).__init__()
        kernel_size = 7
        self.compress = ChannelPool()
        self.spatial = BasicConv(2, 1, kernel_size, stride=1, padding=(kernel_size-1) // 2, relu=False)
    def forward(self, x):
        x_compress = self.compress(x)
        x_out = self.spatial(x_compress)
        scale = F.sigmoid(x_out) # broadcasting
        return x * scale

class CBAM(nn.Module):
    def __init__(self, gate_channels, reduction_ratio=16, pool_types=['avg', 'max'], no_spatial=False):
        super(CBAM, self).__init__()
        self.ChannelGate = ChannelGate(gate_channels, reduction_ratio, pool_types)
        self.no_spatial=no_spatial
        if not no_spatial:
            self.SpatialGate = SpatialGate()
    def forward(self, x):
        x_out = self.ChannelGate(x)
        if not self.no_spatial:
            x_out = self.SpatialGate(x_out)
        return x_out

import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from torch.nn import init


def conv3x3(in_planes, out_planes, stride=1):
    "3x3 convolution with padding"
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
                     padding=1, bias=False)

class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self, inplanes, planes, stride=1, downsample=None, use_cbam=False):
        super(BasicBlock, self).__init__()
        self.conv1 = conv3x3(inplanes, planes, stride)
        self.bn1 = nn.BatchNorm2d(planes)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = conv3x3(planes, planes)
        self.bn2 = nn.BatchNorm2d(planes)
        self.downsample = downsample
        self.stride = stride

        if use_cbam:
            self.cbam = CBAM( planes, 16 )
        else:
            self.cbam = None

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)

        if self.downsample is not None:
            residual = self.downsample(x)

        if not self.cbam is None:
            out = self.cbam(out)

        out += residual
        out = self.relu(out)

        return out

class Bottleneck(nn.Module):
    expansion = 4

    def __init__(self, inplanes, planes, stride=1, downsample=None, use_cbam=False):
        super(Bottleneck, self).__init__()
        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(planes)
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
                               padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(planes)
        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm2d(planes * 4)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride

        if use_cbam:
            self.cbam = CBAM( planes * 4, 16 )
        else:
            self.cbam = None

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)

        out = self.conv3(out)
        out = self.bn3(out)

        if self.downsample is not None:
            residual = self.downsample(x)

        if not self.cbam is None:
            out = self.cbam(out)

        out += residual
        out = self.relu(out)

        return out

class ResNet(nn.Module):
    def __init__(self, block, layers,  network_type, num_classes, att_type=None):
        self.inplanes = 64
        super(ResNet, self).__init__()
        self.network_type = network_type
        # different model config between ImageNet and CIFAR 
        if network_type == "ImageNet":
            self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
            self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
            self.avgpool = nn.AvgPool2d(7)
        else:
            self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)

        self.bn1 = nn.BatchNorm2d(64)
        self.relu = nn.ReLU(inplace=True)

        if att_type=='BAM':
            self.bam1 = BAM(64*block.expansion)
            self.bam2 = BAM(128*block.expansion)
            self.bam3 = BAM(256*block.expansion)
        else:
            self.bam1, self.bam2, self.bam3 = None, None, None

        self.layer1 = self._make_layer(block, 64,  layers[0], att_type=att_type)
        self.layer2 = self._make_layer(block, 128, layers[1], stride=2, att_type=att_type)
        self.layer3 = self._make_layer(block, 256, layers[2], stride=2, att_type=att_type)
        self.layer4 = self._make_layer(block, 512, layers[3], stride=2, att_type=att_type)

        self.fc = nn.Linear(512 * block.expansion, num_classes)

        init.kaiming_normal(self.fc.weight)
        for key in self.state_dict():
            if key.split('.')[-1]=="weight":
                if "conv" in key:
                    init.kaiming_normal(self.state_dict()[key], mode='fan_out')
                if "bn" in key:
                    if "SpatialGate" in key:
                        self.state_dict()[key][...] = 0
                    else:
                        self.state_dict()[key][...] = 1
            elif key.split(".")[-1]=='bias':
                self.state_dict()[key][...] = 0

    def _make_layer(self, block, planes, blocks, stride=1, att_type=None):
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(self.inplanes, planes * block.expansion,
                          kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(planes * block.expansion),
            )

        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample, use_cbam=att_type=='CBAM'))
        self.inplanes = planes * block.expansion
        for i in range(1, blocks):
            layers.append(block(self.inplanes, planes, use_cbam=att_type=='CBAM'))

        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        if self.network_type == "ImageNet":
            x = self.maxpool(x)

        x = self.layer1(x)
        if not self.bam1 is None:
            x = self.bam1(x)

        x = self.layer2(x)
        if not self.bam2 is None:
            x = self.bam2(x)

        x = self.layer3(x)
        if not self.bam3 is None:
            x = self.bam3(x)

        x = self.layer4(x)

        if self.network_type == "ImageNet":
            x = self.avgpool(x)
        else:
            x = F.avg_pool2d(x, 4)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

def ResidualNet(network_type, depth, num_classes, att_type):

    assert network_type in ["ImageNet", "CIFAR10", "CIFAR100"], "network type should be ImageNet or CIFAR10 / CIFAR100"
    assert depth in [18, 34, 50, 101], 'network depth should be 18, 34, 50 or 101'

    if depth == 18:
        model = ResNet(BasicBlock, [2, 2, 2, 2], network_type, num_classes, att_type)

    elif depth == 34:
        model = ResNet(BasicBlock, [3, 4, 6, 3], network_type, num_classes, att_type)

    elif depth == 50:
        model = ResNet(Bottleneck, [3, 4, 6, 3], network_type, num_classes, att_type)

    elif depth == 101:
        model = ResNet(Bottleneck, [3, 4, 23, 3], network_type, num_classes, att_type)

    return model

from google.colab import drive
drive.mount('/content/drive')

os.listdir()

import torch
import math
import torch.nn as nn
import torch.nn.functional as F

class Flatten(nn.Module):
    def forward(self, x):
        return x.view(x.size(0), -1)
class ChannelGate(nn.Module):
    def __init__(self, gate_channel, reduction_ratio=16, num_layers=1):
        super(ChannelGate, self).__init__()
        self.gate_activation = gate_activation
        self.gate_c = nn.Sequential()
        self.gate_c.add_module( 'flatten', Flatten() )
        gate_channels = [gate_channel]
        gate_channels += [gate_channel // reduction_ratio] * num_layers
        gate_channels += [gate_channel]
        for i in range( len(gate_channels) - 2 ):
            self.gate_c.add_module( 'gate_c_fc_%d'%i, nn.Linear(gate_channels[i], gate_channels[i+1]) )
            self.gate_c.add_module( 'gate_c_bn_%d'%(i+1), nn.BatchNorm1d(gate_channels[i+1]) )
            self.gate_c.add_module( 'gate_c_relu_%d'%(i+1), nn.ReLU() )
        self.gate_c.add_module( 'gate_c_fc_final', nn.Linear(gate_channels[-2], gate_channels[-1]) )
    def forward(self, in_tensor):
        avg_pool = F.avg_pool2d( in_tensor, in_tensor.size(2), stride=in_tensor.size(2) )
        return self.gate_c( avg_pool ).unsqueeze(2).unsqueeze(3).expand_as(in_tensor)

class SpatialGate(nn.Module):
    def __init__(self, gate_channel, reduction_ratio=16, dilation_conv_num=2, dilation_val=4):
        super(SpatialGate, self).__init__()
        self.gate_s = nn.Sequential()
        self.gate_s.add_module( 'gate_s_conv_reduce0', nn.Conv2d(gate_channel, gate_channel//reduction_ratio, kernel_size=1))
        self.gate_s.add_module( 'gate_s_bn_reduce0',	nn.BatchNorm2d(gate_channel//reduction_ratio) )
        self.gate_s.add_module( 'gate_s_relu_reduce0',nn.ReLU() )
        for i in range( dilation_conv_num ):
            self.gate_s.add_module( 'gate_s_conv_di_%d'%i, nn.Conv2d(gate_channel//reduction_ratio, gate_channel//reduction_ratio, kernel_size=3, \
						padding=dilation_val, dilation=dilation_val) )
            self.gate_s.add_module( 'gate_s_bn_di_%d'%i, nn.BatchNorm2d(gate_channel//reduction_ratio) )
            self.gate_s.add_module( 'gate_s_relu_di_%d'%i, nn.ReLU() )
        self.gate_s.add_module( 'gate_s_conv_final', nn.Conv2d(gate_channel//reduction_ratio, 1, kernel_size=1) )
    def forward(self, in_tensor):
        return self.gate_s( in_tensor ).expand_as(in_tensor)
class BAM(nn.Module):
    def __init__(self, gate_channel):
        super(BAM, self).__init__()
        self.channel_att = ChannelGate(gate_channel)
        self.spatial_att = SpatialGate(gate_channel)
    def forward(self,in_tensor):
        att = 1 + F.sigmoid( self.channel_att(in_tensor) * self.spatial_att(in_tensor) )
        return att * in_tensor
import torch
import math
import torch.nn as nn
import torch.nn.functional as F

class BasicConv(nn.Module):
    def __init__(self, in_planes, out_planes, kernel_size, stride=1, padding=0, dilation=1, groups=1, relu=True, bn=True, bias=False):
        super(BasicConv, self).__init__()
        self.out_channels = out_planes
        self.conv = nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)
        self.bn = nn.BatchNorm2d(out_planes,eps=1e-5, momentum=0.01, affine=True) if bn else None
        self.relu = nn.ReLU() if relu else None

    def forward(self, x):
        x = self.conv(x)
        if self.bn is not None:
            x = self.bn(x)
        if self.relu is not None:
            x = self.relu(x)
        return x

class Flatten(nn.Module):
    def forward(self, x):
        return x.view(x.size(0), -1)

class ChannelGate(nn.Module):
    def __init__(self, gate_channels, reduction_ratio=16, pool_types=['avg', 'max']):
        super(ChannelGate, self).__init__()
        self.gate_channels = gate_channels
        self.mlp = nn.Sequential(
            Flatten(),
            nn.Linear(gate_channels, gate_channels // reduction_ratio),
            nn.ReLU(),
            nn.Linear(gate_channels // reduction_ratio, gate_channels)
            )
        self.pool_types = pool_types
    def forward(self, x):
        channel_att_sum = None
        for pool_type in self.pool_types:
            if pool_type=='avg':
                avg_pool = F.avg_pool2d( x, (x.size(2), x.size(3)), stride=(x.size(2), x.size(3)))
                channel_att_raw = self.mlp( avg_pool )
            elif pool_type=='max':
                max_pool = F.max_pool2d( x, (x.size(2), x.size(3)), stride=(x.size(2), x.size(3)))
                channel_att_raw = self.mlp( max_pool )
            elif pool_type=='lp':
                lp_pool = F.lp_pool2d( x, 2, (x.size(2), x.size(3)), stride=(x.size(2), x.size(3)))
                channel_att_raw = self.mlp( lp_pool )
            elif pool_type=='lse':
                # LSE pool only
                lse_pool = logsumexp_2d(x)
                channel_att_raw = self.mlp( lse_pool )

            if channel_att_sum is None:
                channel_att_sum = channel_att_raw
            else:
                channel_att_sum = channel_att_sum + channel_att_raw

        scale = F.sigmoid( channel_att_sum ).unsqueeze(2).unsqueeze(3).expand_as(x)
        return x * scale

def logsumexp_2d(tensor):
    tensor_flatten = tensor.view(tensor.size(0), tensor.size(1), -1)
    s, _ = torch.max(tensor_flatten, dim=2, keepdim=True)
    outputs = s + (tensor_flatten - s).exp().sum(dim=2, keepdim=True).log()
    return outputs

class ChannelPool(nn.Module):
    def forward(self, x):
        return torch.cat( (torch.max(x,1)[0].unsqueeze(1), torch.mean(x,1).unsqueeze(1)), dim=1 )

class SpatialGate(nn.Module):
    def __init__(self):
        super(SpatialGate, self).__init__()
        kernel_size = 7
        self.compress = ChannelPool()
        self.spatial = BasicConv(2, 1, kernel_size, stride=1, padding=(kernel_size-1) // 2, relu=False)
    def forward(self, x):
        x_compress = self.compress(x)
        x_out = self.spatial(x_compress)
        scale = F.sigmoid(x_out) # broadcasting
        return x * scale

class CBAM(nn.Module):
    def __init__(self, gate_channels, reduction_ratio=16, pool_types=['avg', 'max'], no_spatial=False):
        super(CBAM, self).__init__()
        self.ChannelGate = ChannelGate(gate_channels, reduction_ratio, pool_types)
        self.no_spatial=no_spatial
        if not no_spatial:
            self.SpatialGate = SpatialGate()
    def forward(self, x):
        x_out = self.ChannelGate(x)
        if not self.no_spatial:
            x_out = self.SpatialGate(x_out)
        return x_out

import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from torch.nn import init


def conv3x3(in_planes, out_planes, stride=1):
    "3x3 convolution with padding"
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
                     padding=1, bias=False)

class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self, inplanes, planes, stride=1, downsample=None, use_cbam=False):
        super(BasicBlock, self).__init__()
        self.conv1 = conv3x3(inplanes, planes, stride)
        self.bn1 = nn.BatchNorm2d(planes)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = conv3x3(planes, planes)
        self.bn2 = nn.BatchNorm2d(planes)
        self.downsample = downsample
        self.stride = stride

        if use_cbam:
            self.cbam = CBAM( planes, 16 )
        else:
            self.cbam = None

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)

        if self.downsample is not None:
            residual = self.downsample(x)

        if not self.cbam is None:
            out = self.cbam(out)

        out += residual
        out = self.relu(out)

        return out

class Bottleneck(nn.Module):
    expansion = 4

    def __init__(self, inplanes, planes, stride=1, downsample=None, use_cbam=False):
        super(Bottleneck, self).__init__()
        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(planes)
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
                               padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(planes)
        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm2d(planes * 4)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride

        if use_cbam:
            self.cbam = CBAM( planes * 4, 16 )
        else:
            self.cbam = None

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)

        out = self.conv3(out)
        out = self.bn3(out)

        if self.downsample is not None:
            residual = self.downsample(x)

        if not self.cbam is None:
            out = self.cbam(out)

        out += residual
        out = self.relu(out)

        return out

class ResNet(nn.Module):
    def __init__(self, block, layers,  network_type, num_classes, att_type=None):
        self.inplanes = 64
        super(ResNet, self).__init__()
        self.network_type = network_type
        # different model config between ImageNet and CIFAR 
        if network_type == "ImageNet":
            self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
            self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
            self.avgpool = nn.AvgPool2d(7)
        else:
            self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)

        self.bn1 = nn.BatchNorm2d(64)
        self.relu = nn.ReLU(inplace=True)

        if att_type=='BAM':
            self.bam1 = BAM(64*block.expansion)
            self.bam2 = BAM(128*block.expansion)
            self.bam3 = BAM(256*block.expansion)
        else:
            self.bam1, self.bam2, self.bam3 = None, None, None

        self.layer1 = self._make_layer(block, 64,  layers[0], att_type=att_type)
        self.layer2 = self._make_layer(block, 128, layers[1], stride=2, att_type=att_type)
        self.layer3 = self._make_layer(block, 256, layers[2], stride=2, att_type=att_type)
        self.layer4 = self._make_layer(block, 512, layers[3], stride=2, att_type=att_type)

        self.fc = nn.Linear(512 * block.expansion, num_classes)

        init.kaiming_normal(self.fc.weight)
        for key in self.state_dict():
            if key.split('.')[-1]=="weight":
                if "conv" in key:
                    init.kaiming_normal(self.state_dict()[key], mode='fan_out')
                if "bn" in key:
                    if "SpatialGate" in key:
                        self.state_dict()[key][...] = 0
                    else:
                        self.state_dict()[key][...] = 1
            elif key.split(".")[-1]=='bias':
                self.state_dict()[key][...] = 0

    def _make_layer(self, block, planes, blocks, stride=1, att_type=None):
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(self.inplanes, planes * block.expansion,
                          kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(planes * block.expansion),
            )

        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample, use_cbam=att_type=='CBAM'))
        self.inplanes = planes * block.expansion
        for i in range(1, blocks):
            layers.append(block(self.inplanes, planes, use_cbam=att_type=='CBAM'))

        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        if self.network_type == "ImageNet":
            x = self.maxpool(x)

        x = self.layer1(x)
        if not self.bam1 is None:
            x = self.bam1(x)

        x = self.layer2(x)
        if not self.bam2 is None:
            x = self.bam2(x)

        x = self.layer3(x)
        if not self.bam3 is None:
            x = self.bam3(x)

        x = self.layer4(x)

        if self.network_type == "ImageNet":
            x = self.avgpool(x)
        else:
            x = F.avg_pool2d(x, 4)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

def ResidualNet(network_type, depth, num_classes, att_type):

    assert network_type in ["ImageNet", "CIFAR10", "CIFAR100"], "network type should be ImageNet or CIFAR10 / CIFAR100"
    assert depth in [18, 34, 50, 101], 'network depth should be 18, 34, 50 or 101'

    if depth == 18:
        model = ResNet(BasicBlock, [2, 2, 2, 2], network_type, num_classes, att_type)

    elif depth == 34:
        model = ResNet(BasicBlock, [3, 4, 6, 3], network_type, num_classes, att_type)

    elif depth == 50:
        model = ResNet(Bottleneck, [3, 4, 6, 3], network_type, num_classes, att_type)

    elif depth == 101:
        model = ResNet(Bottleneck, [3, 4, 23, 3], network_type, num_classes, att_type)

    return model

ResidualNet('CIFAR100', 34, 4, 'cbam')

# # Pre-trained weights: https://github.com/facebookresearch/WSL-Images
# # Apex: https://github.com/NVIDIA/apex
# # Borrowed a lot from abhishek, so give him an upvote: https://www.kaggle.com/abhishek/very-simple-pytorch-training-0-59

# # Parameters

# lr = 2e-5
# img_size = 224
# batch_size = 32
# n_epochs = 10
# n_freeze = 1
# classes = (0, 1, 2, 3)
# coef = [0.5, 1.5, 2.5, 3.5]

    
# # Libraries

# import torch
# import os
# import gc
# import sys
# import cv2
# import subprocess
# import numpy as np
# import pandas as pd
# from tqdm import tqdm
# from albumentations import Compose, RandomBrightnessContrast, ShiftScaleRotate
# from albumentations.pytorch import ToTensor
# from torch.utils.data import Dataset
# from torchvision.models.resnet import ResNet, Bottleneck
# import torch.optim as optim

# # Install Apex for mixed precision

# print('Starting Apex installation ...')

# FNULL = open(os.devnull, 'w')
# process = subprocess.Popen(
#     'pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ../input/nvidia-apex/apex/apex',
#     shell=True, 
#     stdout=FNULL, stderr=subprocess.STDOUT)
# process.wait()

# if process.returncode==0:
#     print('Apex successfully installed')
    
# from apex import amp

# # Functions


# class RetinopathyDatasetTrain(Dataset):
#     def __init__(self, csv_file, transform=None):
#         self.data = pd.read_csv(csv_file)
#         self.transform = transform

#     def __len__(self):
#         return len(self.data)

#     def __getitem__(self, idx):
#         img_name = os.path.join('../input/aptospreprocessed224x/train_images_processed_224x/train_images_processed_224x/', self.data.loc[idx, 'id_code'] + '.png.png') # typo
#         im = cv2.imread(img_name)
#         label = torch.tensor(self.data.loc[idx, 'diagnosis'])
#         if self.transform:
#             augmented = self.transform(image=im)
#             im = augmented['image']
#         return {'image': im, 'labels': label}
    
# class RetinopathyDatasetTest(Dataset):
#     def __init__(self, csv_file, transform=None):
#         self.data = pd.read_csv(csv_file)
#         self.transform = transform

#     def __len__(self):
#         return len(self.data)

#     def __getitem__(self, idx):
#         img_name = os.path.join('../input/aptos2019-blindness-detection/test_images', self.data.loc[idx, 'id_code'] + '.png')
#         im = cv2.imread(img_name)
#         im = circle_crop(im)
#         im = cv2.resize(im, (img_size, img_size))
#         if self.transform:
#             augmented = self.transform(image=im)
#             im = augmented['image']
#         return {'image': im}
    
# def _resnext(path, block, layers, pretrained, progress, **kwargs):
#     model = ResNet(block, layers, **kwargs)
#     model.load_state_dict(torch.load(path))
#     return model

# def resnext101_32x16d_wsl(path, progress=True, **kwargs):
#     """Constructs a ResNeXt-101 32x16 model pre-trained on weakly-supervised data
#     and finetuned on ImageNet from Figure 5 in
#     `"Exploring the Limits of Weakly Supervised Pretraining" <https://arxiv.org/abs/1805.00932>`_
#     Args:
#         progress (bool): If True, displays a progress bar of the download to stderr.
#     """
#     kwargs['groups'] = 32
#     kwargs['width_per_group'] = 16
#     return _resnext(path, Bottleneck, [3, 4, 23, 3], True, progress, **kwargs)

# def train_model(model,  n_epochs, classification=True):  
    
#     for epoch in range(n_epochs):         
#         if epoch == n_freeze:                
#             for param in model.parameters():
#                 param.requires_grad = True                
#         tr_loss = 0        
#         counter = 0
#         model.train()
#         print('Epoch {}/{}'.format(epoch, n_epochs - 1))
#         print('-' * 10)        
#         for step, batch in enumerate(train_data_loader):
#             if classification:
#                 inputs = batch["image"]
#                 labels = batch["labels"]
#                 inputs = inputs.to(device, dtype=torch.float)
#                 labels = labels.to(device, dtype=torch.long)
#             else:
#                 inputs = batch["image"]
#                 labels = batch["labels"].view(-1, 1)
#                 inputs = inputs.to(device, dtype=torch.float)
#                 labels = labels.to(device, dtype=torch.float)
#             outputs = model(inputs)
#             loss = criterion(outputs, labels)
#             with amp.scale_loss(loss, optimizer) as scaled_loss:
#                 scaled_loss.backward()     
#             tr_loss += loss.item()            
#             optimizer.step()
#             optimizer.zero_grad()
#         epoch_loss = tr_loss / len(train_data_loader)
#         print('Training Loss: {:.4f}'.format(epoch_loss))
        
#     return model

# # processing

# transform_train = Compose([
#     ShiftScaleRotate(
#         shift_limit=0.1,
#         scale_limit=0.1,
#         rotate_limit=365,
#         p=1.0),
#     RandomBrightnessContrast(p=1.0),
#     ToTensor()
# ])

# transform_test = Compose([
#     ToTensor()
# ])

# train_dataset = RetinopathyDatasetTrain(csv_file='../input/aptos2019-blindness-detection/train.csv', transform=transform_train)
# train_data_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)
    
# test_dataset = RetinopathyDatasetTest(csv_file='../input/aptos2019-blindness-detection/test.csv', transform=transform_test)
# test_data_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=4)
    
# device = torch.device("cuda:0")

# # classification model
    
# model = resnext101_32x16d_wsl(path='../input/ig-resnext101-32x16/ig_resnext101_32x16-c6f796b0.pth')

# for param in model.parameters():
#     param.requires_grad = False

# model.fc = torch.nn.Linear(2048, len(classes))
# model.to(device)

# criterion = torch.nn.CrossEntropyLoss()
# plist = [{'params': model.parameters(), 'lr': lr}]
# optimizer = optim.Adam(plist, lr=lr)

# model, optimizer = amp.initialize(model, optimizer, opt_level="O1")

# model = train_model(model=model,n_epochs=n_epochs,classification=True)

# for param in model.parameters():
#     param.requires_grad = False

# model.eval()

# test_preds1 = np.zeros((len(test_dataset), 5))

# for i, x_batch in enumerate(tqdm(test_data_loader)):
#     x_batch = x_batch["image"]
#     pred = model(x_batch.to(device))
#     test_preds1[i * batch_size:(i + 1) * batch_size] = pred.detach().cpu()

# test_preds1 = np.argmax(test_preds1, axis=1)  

# del(model, optimizer, criterion, plist)
# gc.collect()
# torch.cuda.empty_cache()

# # regression model

# model = resnext101_32x16d_wsl(path='../input/ig-resnext101-32x16/ig_resnext101_32x16-c6f796b0.pth')

# for param in model.parameters():
#     param.requires_grad = False

# model.fc = torch.nn.Linear(2048, 1)
# model.to(device)

# criterion = torch.nn.MSELoss()
# plist = [{'params': model.parameters(), 'lr': lr}]
# optimizer = optim.Adam(plist, lr=lr)

# model, optimizer = amp.initialize(model, optimizer, opt_level="O1")

# model = train_model(model=model,n_epochs=n_epochs,classification=False)

# for param in model.parameters():
#     param.requires_grad = False

# model.eval()

# test_preds2 = np.zeros((len(test_dataset), 1))

# for i, x_batch in enumerate(tqdm(test_data_loader)):
#     x_batch = x_batch["image"]
#     pred = model(x_batch.to(device))
#     test_preds2[i * batch_size:(i + 1) * batch_size] = pred.detach().cpu()

# i = 0

# for pred in test_preds2:
#     if pred < coef[0]:
#         test_preds2[i] = 0
#         i += 1
#     elif pred >= coef[0] and pred < coef[1]:
#         test_preds2[i] = 1
#         i += 1
#     elif pred >= coef[1] and pred < coef[2]:
#         test_preds2[i] = 2
#         i += 1
#     elif pred >= coef[2] and pred < coef[3]:
#         test_preds2[i] = 3
#         i += 1
#     else:
#         test_preds2[i] = 4
#         i += 1

# test_preds2 = test_preds2.reshape(-1,)        
        
# # combine

# final_pred = np.round((test_preds1+test_preds2)/2)
        
# # submit
 
# sample = pd.read_csv("../input/aptos2019-blindness-detection/sample_submission.csv")
# sample.diagnosis = final_pred.astype(int)
# sample.to_csv("submission.csv", index=False)