run resnet & FAA getBraTS_5

+import os
+import time
+import importlib
+import collections
+import pickle as cp
+import numpy as np
+
+
+import torch
+import torchvision
+import torch.nn.functional as F
+import torchvision.models as models
+import torchvision.transforms as transforms
+from torch.utils.data import Subset
+
+from sklearn.model_selection import StratifiedShuffleSplit
+
+from networks import basenet
+from networks import grayResNet
+
+
+DATASET_PATH = './data/'
+current_epoch = 0
+
+
+def split_dataset(args, dataset, k):
+    # load dataset
+    X = list(range(len(dataset)))
+    Y = dataset.targets
+
+    # split to k-fold
+    assert len(X) == len(Y)
+
+    def _it_to_list(_it):
+        return list(zip(*list(_it)))
+
+    sss = StratifiedShuffleSplit(n_splits=k, random_state=args.seed, test_size=0.1)
+    Dm_indexes, Da_indexes = _it_to_list(sss.split(X, Y))
+
+    return Dm_indexes, Da_indexes
+
+
+def concat_image_features(image, features, max_features=3):
+    _, h, w = image.shape
+
+    max_features = min(features.size(0), max_features)
+    image_feature = image.clone()
+
+    for i in range(max_features):
+        feature = features[i:i+1]
+        _min, _max = torch.min(feature), torch.max(feature)
+        feature = (feature - _min) / (_max - _min + 1e-6)
+        feature = torch.cat([feature]*3, 0)
+        feature = feature.view(1, 3, feature.size(1), feature.size(2))
+        feature = F.upsample(feature, size=(h,w), mode="bilinear")
+        feature = feature.view(3, h, w)
+        image_feature = torch.cat((image_feature, feature), 2)
+
+    return image_feature
+
+
+def get_model_name(args):
+    from datetime import datetime
+    now = datetime.now()
+    date_time = now.strftime("%B_%d_%H:%M:%S")
+    model_name = '__'.join([date_time, args.network, str(args.seed)])
+    return model_name
+
+
+def dict_to_namedtuple(d):
+    Args = collections.namedtuple('Args', sorted(d.keys()))
+
+    for k,v in d.items():
+        if type(v) is dict:
+            d[k] = dict_to_namedtuple(v)
+
+        elif type(v) is str:
+            try:
+                d[k] = eval(v)
+            except:
+                d[k] = v
+
+    args = Args(**d)
+    return args
+
+
+def parse_args(kwargs):
+    # combine with default args
+    kwargs['dataset'] =  kwargs['dataset'] if 'dataset' in kwargs else 'cifar10'
+    kwargs['network'] =  kwargs['network'] if 'network' in kwargs else 'resnet_cifar10'
+    kwargs['optimizer'] = kwargs['optimizer'] if 'optimizer' in kwargs else 'adam'
+    kwargs['learning_rate'] = kwargs['learning_rate'] if 'learning_rate' in kwargs else 0.1
+    kwargs['seed'] =  kwargs['seed'] if 'seed' in kwargs else None
+    kwargs['use_cuda'] =  kwargs['use_cuda'] if 'use_cuda' in kwargs else True
+    kwargs['use_cuda'] =  kwargs['use_cuda'] and torch.cuda.is_available()
+    kwargs['num_workers'] = kwargs['num_workers'] if 'num_workers' in kwargs else 4
+    kwargs['print_step'] = kwargs['print_step'] if 'print_step' in kwargs else 2000
+    kwargs['val_step'] = kwargs['val_step'] if 'val_step' in kwargs else 2000
+    kwargs['scheduler'] = kwargs['scheduler'] if 'scheduler' in kwargs else 'exp'
+    kwargs['batch_size'] = kwargs['batch_size'] if 'batch_size' in kwargs else 128
+    kwargs['start_step'] = kwargs['start_step'] if 'start_step' in kwargs else 0
+    kwargs['max_step'] = kwargs['max_step'] if 'max_step' in kwargs else 64000
+    kwargs['fast_auto_augment'] = kwargs['fast_auto_augment'] if 'fast_auto_augment' in kwargs else False
+    kwargs['augment_path'] = kwargs['augment_path'] if 'augment_path' in kwargs else None
+
+    # to named tuple
+    args = dict_to_namedtuple(kwargs)
+    return args, kwargs
+
+
+def select_model(args):
+    resnet_dict = {'ResNet18':grayResNet.ResNet18(), 'ResNet34':grayResNet.ResNet34(), 
+    'ResNet50':grayResNet.ResNet50(), 'ResNet101':grayResNet.ResNet101(), 'ResNet152':grayResNet.ResNet152()}
+    #print("args.network: \n", args.network)
+    if args.network in resnet_dict:
+        backbone = resnet_dict[args.network]
+        model = basenet.BaseNet(backbone, args)
+    else:
+        Net = getattr(importlib.import_module('networks.{}'.format(args.network)), 'Net')
+        model = Net(args)
+
+    print(model)
+    return model
+
+
+def select_optimizer(args, model):
+    if args.optimizer == 'sgd':
+        optimizer = torch.optim.SGD(model.parameters(), lr=args.learning_rate, momentum=0.9, weight_decay=0.0001)
+    elif args.optimizer == 'rms':
+        #optimizer = torch.optim.RMSprop(model.parameters(), lr=args.learning_rate, momentum=0.9, weight_decay=1e-5)
+        optimizer = torch.optim.RMSprop(model.parameters(), lr=args.learning_rate)
+    elif args.optimizer == 'adam':
+        optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate)
+    else:
+        raise Exception('Unknown Optimizer')
+    return optimizer
+
+
+def select_scheduler(args, optimizer):
+    if not args.scheduler or args.scheduler == 'None':
+        return None
+    elif args.scheduler =='clr':
+        return torch.optim.lr_scheduler.CyclicLR(optimizer, 0.01, 0.015, mode='triangular2', step_size_up=250000, cycle_momentum=False)
+    elif args.scheduler =='exp':
+        return torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.9999283, last_epoch=-1)
+    else:
+        raise Exception('Unknown Scheduler')
+
+
+def get_dataset(args, transform, split='train'):
+    assert split in ['train', 'val', 'test', 'trainval']
+
+    if args.dataset == 'cifar10':
+        train = split in ['train', 'val', 'trainval']
+        dataset = torchvision.datasets.CIFAR10(DATASET_PATH,
+                                               train=train,
+                                               transform=transform,
+                                               download=True)
+
+        if split in ['train', 'val']:
+            split_path = os.path.join(DATASET_PATH,
+                    'cifar-10-batches-py', 'train_val_index.cp')
+
+            if not os.path.exists(split_path):
+                [train_index], [val_index] = split_dataset(args, dataset, k=1)
+                split_index = {'train':train_index, 'val':val_index}
+                cp.dump(split_index, open(split_path, 'wb'))
+
+            split_index = cp.load(open(split_path, 'rb'))
+            dataset = Subset(dataset, split_index[split])
+
+    elif args.dataset == 'imagenet':
+        dataset = torchvision.datasets.ImageNet(DATASET_PATH,
+                                                split=split,
+                                                transform=transform,
+                                                download=(split is 'val'))
+
+    else:
+        raise Exception('Unknown dataset')
+
+    return dataset
+
+
+def get_dataloader(args, dataset, shuffle=False, pin_memory=True):
+    data_loader = torch.utils.data.DataLoader(dataset,
+                                              batch_size=args.batch_size,
+                                              shuffle=shuffle,
+                                              num_workers=args.num_workers,
+                                              pin_memory=pin_memory)
+    return data_loader
+
+
+def get_inf_dataloader(args, dataset):
+    global current_epoch
+    data_loader = iter(get_dataloader(args, dataset, shuffle=True))
+
+    while True:
+        try:
+            batch = next(data_loader)
+
+        except StopIteration:
+            current_epoch += 1
+            data_loader = iter(get_dataloader(args, dataset, shuffle=True))
+            batch = next(data_loader)
+
+        yield batch
+
+
+def get_train_transform(args, model, log_dir=None):
+    if args.fast_auto_augment:
+        assert args.dataset == 'cifar10' # TODO: FastAutoAugment for Imagenet
+
+        from fast_auto_augment import fast_auto_augment
+        if args.augment_path:
+            transform = cp.load(open(args.augment_path, 'rb'))
+            os.system('cp {} {}'.format(
+                args.augment_path, os.path.join(log_dir, 'augmentation.cp')))
+        else:
+            transform = fast_auto_augment(args, model, K=4, B=1, num_process=4)
+            if log_dir:
+                cp.dump(transform, open(os.path.join(log_dir, 'augmentation.cp'), 'wb'))
+
+    elif args.dataset == 'cifar10':
+        transform = transforms.Compose([
+            transforms.Pad(4),
+            transforms.RandomCrop(32),
+            transforms.RandomHorizontalFlip(),
+            transforms.ToTensor()
+        ])
+
+    elif args.dataset == 'imagenet':
+        resize_h, resize_w = model.img_size[0], int(model.img_size[1]*1.875)
+        transform = transforms.Compose([
+            transforms.Resize([resize_h, resize_w]),
+            transforms.RandomCrop(model.img_size),
+            transforms.RandomHorizontalFlip(),
+            transforms.ToTensor()
+        ])
+
+    else:
+        raise Exception('Unknown Dataset')
+
+    print(transform)
+
+    return transform
+
+
+def get_valid_transform(args, model):
+    if args.dataset == 'cifar10':
+        val_transform = transforms.Compose([
+            transforms.Resize(32),
+            transforms.ToTensor()
+        ])
+
+    elif args.dataset == 'imagenet':
+        resize_h, resize_w = model.img_size[0], int(model.img_size[1]*1.875)
+        val_transform = transforms.Compose([
+            transforms.Resize([resize_h, resize_w]),
+            transforms.ToTensor()
+        ])
+
+    else:
+        raise Exception('Unknown Dataset')
+
+    return val_transform
+
+
+def train_step(args, model, optimizer, scheduler, criterion, batch, step, writer, device=None):
+    model.train()
+    images, target = batch
+
+    if device:
+        images = images.to(device)
+        target = target.to(device)
+
+    elif args.use_cuda:
+        images = images.cuda(non_blocking=True)
+        target = target.cuda(non_blocking=True)
+
+    # compute output
+    start_t = time.time()
+    output, first = model(images)
+    forward_t = time.time() - start_t
+    loss = criterion(output, target)
+
+    # measure accuracy and record loss
+    acc1, acc5 = accuracy(output, target, topk=(1, 5))
+    acc1 /= images.size(0)
+    acc5 /= images.size(0)
+
+    # compute gradient and do SGD step
+    optimizer.zero_grad()
+    start_t = time.time()
+    loss.backward()
+    backward_t = time.time() - start_t
+    optimizer.step()
+    if scheduler: scheduler.step()
+
+    if writer and step % args.print_step == 0:
+        n_imgs = min(images.size(0), 10)
+        for j in range(n_imgs):
+            writer.add_image('train/input_image',
+                    concat_image_features(images[j], first[j]), global_step=step)
+
+    return acc1, acc5, loss, forward_t, backward_t
+
+
+def validate(args, model, criterion, valid_loader, step, writer, device=None):
+    # switch to evaluate mode
+    model.eval()
+
+    acc1, acc5 = 0, 0
+    samples = 0
+    infer_t = 0
+
+    with torch.no_grad():
+        for i, (images, target) in enumerate(valid_loader):
+
+            start_t = time.time()
+            if device:
+                images = images.to(device)
+                target = target.to(device)
+
+            elif args.use_cuda is not None:
+                images = images.cuda(non_blocking=True)
+                target = target.cuda(non_blocking=True)
+
+            # compute output
+            output, first = model(images)
+            loss = criterion(output, target)
+            infer_t += time.time() - start_t
+
+            # measure accuracy and record loss
+            _acc1, _acc5 = accuracy(output, target, topk=(1, 5))
+            acc1 += _acc1
+            acc5 += _acc5
+            samples += images.size(0)
+
+    acc1 /= samples
+    acc5 /= samples
+
+    if writer:
+        n_imgs = min(images.size(0), 10)
+        for j in range(n_imgs):
+            writer.add_image('valid/input_image',
+                    concat_image_features(images[j], first[j]), global_step=step)
+
+    return acc1, acc5, loss, infer_t
+
+
+def accuracy(output, target, topk=(1,)):
+    """Computes the accuracy over the k top predictions for the specified values of k"""
+    with torch.no_grad():
+        maxk = max(topk)
+        batch_size = target.size(0)
+
+        _, pred = output.topk(maxk, 1, True, True)
+        pred = pred.t()
+        correct = pred.eq(target.view(1, -1).expand_as(pred))
+
+        res = []
+        for k in topk:
+            correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
+            res.append(correct_k)
+        return res

@@ -54,10 +54,13 @@ def train_child(args, model, dataset, subset_indx, device=None):
         if torch.cuda.device_count() > 1:
             print('\n[+] Use {} GPUs'.format(torch.cuda.device_count()))
             model = nn.DataParallel(model)
+        elif torch.cuda.device_count() == 1:
+            print('\n[+] Use {} GPUs'.format(torch.cuda.device_count()))
 
     start_t = time.time()
     for step in range(args.start_step, args.max_step):
         batch = next(data_loader)
+
         _train_res = train_step(args, model, optimizer, scheduler, criterion, batch, step, None, device)
 
         if step % args.print_step == 0:
@@ -173,7 +176,7 @@ def process_fn(args_str, model, dataset, Dm_indx, Da_indx, T, transform_candidat
     device = torch.device('cuda:%d' % device_id)
     _transform = []
 
-    print('[+] Child %d training strated (GPU: %d)' % (k, device_id))
+    print('[+] Child %d training started (GPU: %d)' % (k, device_id))
 
     # train child model
     child_model = copy.deepcopy(model)
@@ -188,7 +191,7 @@ def process_fn(args_str, model, dataset, Dm_indx, Da_indx, T, transform_candidat
 
     return _transform
 
-
+#fast_auto_augment(args, model, K=4, B=1, num_process=4)
 def fast_auto_augment(args, model, transform_candidates=None, K=5, B=100, T=2, N=10, num_process=5):
     args_str = json.dumps(args._asdict())
     dataset = get_dataset(args, None, 'trainval')

@@ -4,6 +4,12 @@ class BaseNet(nn.Module):
     def __init__(self, backbone, args):
         super(BaseNet, self).__init__()
 
+        #testing
+        for layer in backbone.children():
+            print("\nRESNET50 LAYERS\n")
+            print(layer)
+
+
         # Separate layers
         self.first = nn.Sequential(*list(backbone.children())[:1])
         self.after = nn.Sequential(*list(backbone.children())[1:-1])
@@ -14,6 +20,20 @@ class BaseNet(nn.Module):
     def forward(self, x):
         f = self.first(x)
         x = self.after(f)
-        x = x.reshape(x.size(0), -1)
         x = self.fc(x)
         return x, f
+
+
+"""
+        print("before reshape:\n", x.size()) 
+        #[128, 2048, 4, 4] 
+        # #cifar 내장[128, 2048, 1, 1]
+        x = x.reshape(x.size(0), -1)
+        print("after reshape:\n",  x.size()) 
+        #[128, 32768] 
+        #cifar [128, 2048]
+        #RuntimeError: size mismatch, m1: [128 x 32768], m2: [2048 x 10] 
+        print("fc :\n", self.fc) 
+        #Linear(in_features=2048, out_features=10, bias=True)
+        #cifar Linear(in_features=2048, out_features=1000, bias=True)
+"""

+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, in_planes, planes, stride=1):
+        super(BasicBlock, self).__init__()
+        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+
+        self.shortcut = nn.Sequential()
+        if stride != 1 or in_planes != self.expansion*planes:
+            self.shortcut = nn.Sequential(
+                nn.Conv2d(in_planes, self.expansion*planes, kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(self.expansion*planes)
+            )
+
+    def forward(self, x):
+        out = F.relu(self.bn1(self.conv1(x)))
+        out = self.bn2(self.conv2(out))
+        out += self.shortcut(x)
+        out = F.relu(out)
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, in_planes, planes, stride=1):
+        super(Bottleneck, self).__init__()
+        self.conv1 = nn.Conv2d(in_planes, 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, self.expansion*planes, kernel_size=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(self.expansion*planes)
+
+        self.shortcut = nn.Sequential()
+        if stride != 1 or in_planes != self.expansion*planes:
+            self.shortcut = nn.Sequential(
+                nn.Conv2d(in_planes, self.expansion*planes, kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(self.expansion*planes)
+            )
+
+    def forward(self, x):
+        out = F.relu(self.bn1(self.conv1(x)))
+        out = F.relu(self.bn2(self.conv2(out)))
+        out = self.bn3(self.conv3(out))
+        out += self.shortcut(x)
+        out = F.relu(out)
+        return out
+
+
+class ResNet(nn.Module):
+    def __init__(self, block, num_blocks, num_classes=10):
+        super(ResNet, self).__init__()
+        self.in_planes = 64
+
+        self.conv1 = nn.Conv2d(1, 64, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1)
+        self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2)
+        self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2)
+        self.linear = nn.Linear(512*block.expansion, num_classes)
+
+    def _make_layer(self, block, planes, num_blocks, stride):
+        strides = [stride] + [1]*(num_blocks-1)
+        layers = []
+        for stride in strides:
+            layers.append(block(self.in_planes, planes, stride))
+            self.in_planes = planes * block.expansion
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        out = F.relu(self.bn1(self.conv1(x)))
+        out = self.layer1(out)
+        out = self.layer2(out)
+        out = self.layer3(out)
+        out = self.layer4(out)
+        out = F.avg_pool2d(out, 4)
+        out = out.view(out.size(0), -1)
+        out = self.linear(out)
+        return out
+
+
+def ResNet18():
+    return ResNet(BasicBlock, [2,2,2,2])
+
+def ResNet34():
+    return ResNet(BasicBlock, [3,4,6,3])
+
+def ResNet50():
+    return ResNet(Bottleneck, [3,4,6,3])
+
+def ResNet101():
+    return ResNet(Bottleneck, [3,4,23,3])
+
+def ResNet152():
+    return ResNet(Bottleneck, [3,8,36,3])
+
+
+def test():
+    net = ResNet18()
+    y = net(torch.randn(1,3,32,32))
+    print(y.size())

+import torch
+import torch.nn as nn
+#from .utils import load_state_dict_from_url
+
+
+__all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',
+           'resnet152', 'resnext50_32x4d', 'resnext101_32x8d',
+           'wide_resnet50_2', 'wide_resnet101_2']
+
+
+model_urls = {
+    'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
+    'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
+    'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
+    'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
+    'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
+    'resnext50_32x4d': 'https://download.pytorch.org/models/resnext50_32x4d-7cdf4587.pth',
+    'resnext101_32x8d': 'https://download.pytorch.org/models/resnext101_32x8d-8ba56ff5.pth',
+    'wide_resnet50_2': 'https://download.pytorch.org/models/wide_resnet50_2-95faca4d.pth',
+    'wide_resnet101_2': 'https://download.pytorch.org/models/wide_resnet101_2-32ee1156.pth',
+}
+
+
+def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=dilation, groups=groups, bias=False, dilation=dilation)
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+    """1x1 convolution"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
+                 base_width=64, dilation=1, norm_layer=None):
+        super(BasicBlock, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        if groups != 1 or base_width != 64:
+            raise ValueError('BasicBlock only supports groups=1 and base_width=64')
+        if dilation > 1:
+            raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
+        # Both self.conv1 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = norm_layer(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = norm_layer(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = 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:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    # Bottleneck in torchvision places the stride for downsampling at 3x3 convolution(self.conv2)
+    # while original implementation places the stride at the first 1x1 convolution(self.conv1)
+    # according to "Deep residual learning for image recognition"https://arxiv.org/abs/1512.03385.
+    # This variant is also known as ResNet V1.5 and improves accuracy according to
+    # https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch.
+
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
+                 base_width=64, dilation=1, norm_layer=None):
+        super(Bottleneck, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        width = int(planes * (base_width / 64.)) * groups
+        # Both self.conv2 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = conv1x1(inplanes, width)
+        self.bn1 = norm_layer(width)
+        self.conv2 = conv3x3(width, width, stride, groups, dilation)
+        self.bn2 = norm_layer(width)
+        self.conv3 = conv1x1(width, planes * self.expansion)
+        self.bn3 = norm_layer(planes * self.expansion)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = 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:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet(nn.Module):
+
+    def __init__(self, block, layers, num_classes=1000, zero_init_residual=False,
+                 groups=1, width_per_group=64, replace_stride_with_dilation=None,
+                 norm_layer=None):
+        super(ResNet, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        self._norm_layer = norm_layer
+
+        self.inplanes = 64
+        self.dilation = 1
+        if replace_stride_with_dilation is None:
+            # each element in the tuple indicates if we should replace
+            # the 2x2 stride with a dilated convolution instead
+            replace_stride_with_dilation = [False, False, False]
+        if len(replace_stride_with_dilation) != 3:
+            raise ValueError("replace_stride_with_dilation should be None "
+                             "or a 3-element tuple, got {}".format(replace_stride_with_dilation))
+        self.groups = groups
+        self.base_width = width_per_group
+        # change dimension 3->1 for grayscale input
+        self.conv1 = nn.Conv2d(1, self.inplanes, kernel_size=7, stride=2, padding=3,
+                               bias=False)
+        self.bn1 = norm_layer(self.inplanes)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2,
+                                       dilate=replace_stride_with_dilation[0])
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=2,
+                                       dilate=replace_stride_with_dilation[1])
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=2,
+                                       dilate=replace_stride_with_dilation[2])
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        self.fc = nn.Linear(512 * block.expansion, num_classes)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+        # Zero-initialize the last BN in each residual branch,
+        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
+        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
+        if zero_init_residual:
+            for m in self.modules():
+                if isinstance(m, Bottleneck):
+                    nn.init.constant_(m.bn3.weight, 0)
+                elif isinstance(m, BasicBlock):
+                    nn.init.constant_(m.bn2.weight, 0)
+
+    def _make_layer(self, block, planes, blocks, stride=1, dilate=False):
+        norm_layer = self._norm_layer
+        downsample = None
+        previous_dilation = self.dilation
+        if dilate:
+            self.dilation *= stride
+            stride = 1
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                conv1x1(self.inplanes, planes * block.expansion, stride),
+                norm_layer(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample, self.groups,
+                            self.base_width, previous_dilation, norm_layer))
+        self.inplanes = planes * block.expansion
+        for _ in range(1, blocks):
+            layers.append(block(self.inplanes, planes, groups=self.groups,
+                                base_width=self.base_width, dilation=self.dilation,
+                                norm_layer=norm_layer))
+
+        return nn.Sequential(*layers)
+
+    def _forward_impl(self, x):
+        # See note [TorchScript super()]
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+
+        x = self.avgpool(x)
+        x = torch.flatten(x, 1)
+        x = self.fc(x)
+
+        return x
+
+    def forward(self, x):
+        return self._forward_impl(x)
+
+
+def _resnet(arch, block, layers, pretrained, progress, **kwargs):
+    model = ResNet(block, layers, **kwargs)
+    # if pretrained:
+    #     state_dict = load_state_dict_from_url(model_urls[arch],
+    #                                           progress=progress)
+    #     model.load_state_dict(state_dict)
+    return model
+
+
+def resnet18(pretrained=False, progress=True, **kwargs):
+    r"""ResNet-18 model from
+    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return _resnet('resnet18', BasicBlock, [2, 2, 2, 2], pretrained, progress,
+                   **kwargs)
+
+
+def resnet34(pretrained=False, progress=True, **kwargs):
+    r"""ResNet-34 model from
+    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return _resnet('resnet34', BasicBlock, [3, 4, 6, 3], pretrained, progress,
+                   **kwargs)
+
+
+def resnet50(pretrained=False, progress=True, **kwargs):
+    r"""ResNet-50 model from
+    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return _resnet('resnet50', Bottleneck, [3, 4, 6, 3], pretrained, progress,
+                   **kwargs)
+
+
+def resnet101(pretrained=False, progress=True, **kwargs):
+    r"""ResNet-101 model from
+    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return _resnet('resnet101', Bottleneck, [3, 4, 23, 3], pretrained, progress,
+                   **kwargs)
+
+
+def resnet152(pretrained=False, progress=True, **kwargs):
+    r"""ResNet-152 model from
+    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return _resnet('resnet152', Bottleneck, [3, 8, 36, 3], pretrained, progress,
+                   **kwargs)
+
+
+def resnext50_32x4d(pretrained=False, progress=True, **kwargs):
+    r"""ResNeXt-50 32x4d model from
+    `"Aggregated Residual Transformation for Deep Neural Networks" <https://arxiv.org/pdf/1611.05431.pdf>`_
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    kwargs['groups'] = 32
+    kwargs['width_per_group'] = 4
+    return _resnet('resnext50_32x4d', Bottleneck, [3, 4, 6, 3],
+                   pretrained, progress, **kwargs)
+
+
+def resnext101_32x8d(pretrained=False, progress=True, **kwargs):
+    r"""ResNeXt-101 32x8d model from
+    `"Aggregated Residual Transformation for Deep Neural Networks" <https://arxiv.org/pdf/1611.05431.pdf>`_
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    kwargs['groups'] = 32
+    kwargs['width_per_group'] = 8
+    return _resnet('resnext101_32x8d', Bottleneck, [3, 4, 23, 3],
+                   pretrained, progress, **kwargs)
+
+
+def wide_resnet50_2(pretrained=False, progress=True, **kwargs):
+    r"""Wide ResNet-50-2 model from
+    `"Wide Residual Networks" <https://arxiv.org/pdf/1605.07146.pdf>`_
+    The model is the same as ResNet except for the bottleneck number of channels
+    which is twice larger in every block. The number of channels in outer 1x1
+    convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048
+    channels, and in Wide ResNet-50-2 has 2048-1024-2048.
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    kwargs['width_per_group'] = 64 * 2
+    return _resnet('wide_resnet50_2', Bottleneck, [3, 4, 6, 3],
+                   pretrained, progress, **kwargs)
+
+
+def wide_resnet101_2(pretrained=False, progress=True, **kwargs):
+    r"""Wide ResNet-101-2 model from
+    `"Wide Residual Networks" <https://arxiv.org/pdf/1605.07146.pdf>`_
+    The model is the same as ResNet except for the bottleneck number of channels
+    which is twice larger in every block. The number of channels in outer 1x1
+    convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048
+    channels, and in Wide ResNet-50-2 has 2048-1024-2048.
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    kwargs['width_per_group'] = 64 * 2
+    return _resnet('wide_resnet101_2', Bottleneck, [3, 4, 23, 3],
+                   pretrained, progress, **kwargs)

@@ -6,6 +6,7 @@ import pickle as cp
 import glob
 import numpy as np
 import pandas as pd
+
 from natsort import natsorted
 from PIL import Image
 import torch
@@ -21,6 +22,7 @@ from sklearn.model_selection import train_test_split
 from sklearn.model_selection import KFold
 
 from networks import basenet
+from networks import grayResNet, grayResNet2
 
 DATASET_PATH = '/content/drive/My Drive/CD2 Project/data/BraTS_Training/train_frame/'
 TRAIN_DATASET_PATH = '/content/drive/My Drive/CD2 Project/data/BraTS_Training/train_frame/'
@@ -55,40 +57,6 @@ def split_dataset(args, dataset, k):
 
     return Dm_indexes, Da_indexes
 
-def split_dataset2222(args, dataset, k):
-    # load dataset
-    X = list(range(len(dataset)))
-    
-    # split to k-fold
-    #assert len(X) == len(Y)
-
-    def _it_to_list(_it):
-        return list(zip(*list(_it)))
-
-    x_train = ()
-    x_test = ()
-
-    for i in range(k):
-        #xtr, xte = train_test_split(X, random_state=args.seed, test_size=0.1)
-        xtr, xte = train_test_split(X, random_state=None, test_size=0.1)
-        x_train.append(np.array(xtr))
-        x_test.append(np.array(xte))
-
-    y_train = np.array([0]* len(x_train))
-    y_test = np.array([0]* len(x_test))
-
-    x_train = tuple(x_train)
-    x_test = tuple(x_test)
-
-    trainset = (zip(x_train, y_train),)
-    testset = (zip(x_test, y_test),)
-
-    Dm_indexes, Da_indexes = trainset, testset
-
-    print(type(Dm_indexes), np.shape(Dm_indexes))
-    print("DM\n", np.shape(Dm_indexes), Dm_indexes, "\nDA\n", np.shape(Da_indexes), Da_indexes)
-
-    return Dm_indexes, Da_indexes
 
 def concat_image_features(image, features, max_features=3):
     _, h, w = image.shape
@@ -159,8 +127,22 @@ def parse_args(kwargs):
 
 
 def select_model(args):
-    if args.network in models.__dict__:
-        backbone = models.__dict__[args.network]()
+    # resnet_dict = {'ResNet18':grayResNet.ResNet18(), 'ResNet34':grayResNet.ResNet34(), 
+    # 'ResNet50':grayResNet.ResNet50(), 'ResNet101':grayResNet.ResNet101(), 'ResNet152':grayResNet.ResNet152()}
+    
+
+    # grayResNet2
+    resnet_dict = {'resnet18':grayResNet2.resnet18(), 'resnet34':grayResNet2.resnet34(), 
+    'resnet50':grayResNet2.resnet50(), 'resnet101':grayResNet2.resnet101(), 'resnet152':grayResNet2.resnet152()}
+
+    if args.network in resnet_dict:
+        backbone = resnet_dict[args.network]
+        #testing
+        # print("\nRESNET50 LAYERS\n")
+        # for layer in backbone.children():
+        #     print(layer)
+        # print("LAYER THE END\n")
+
         model = basenet.BaseNet(backbone, args)
     else:
         Net = getattr(importlib.import_module('networks.{}'.format(args.network)), 'Net')