Merge pull request #1 from chenjiayun212/dev-xyc

add cosoc

Author: chenjiayun212

config/classifiers/COSOC.yaml

@@ -1,3 +1,8 @@
 classifier:
   name: COSOC
-  kwargs: ~
+  kwargs: 
+    alpha: 0.8
+    beta: 0.8
+    num_patches: 7
+    fsl_alg: CC
+    

core/data/collates/collate_functions.py

@@ -156,18 +156,20 @@ def method(self, batch):
             # global_labels = torch.tensor(labels,dtype=torch.int64)
             # global_labels = torch.tensor(labels,dtype=torch.int64).reshape(self.episode_size,self.way_num,
             # self.shot_num*self.times+self.query_num)
+            patch_mode = True
             global_labels = torch.tensor(labels, dtype=torch.int64).reshape(
                 -1, self.way_num, self.shot_num + self.query_num
             )
-            global_labels = (
-                global_labels[..., 0]
-                .unsqueeze(-1)
-                .repeat(
-                    1,
-                    1,
-                    self.shot_num * self.times + self.query_num * self.times_q,
+            if not patch_mode:
+                global_labels = (
+                    global_labels[..., 0]
+                    .unsqueeze(-1)
+                    .repeat(
+                        1,
+                        1,
+                        self.shot_num * self.times + self.query_num * self.times_q,
+                    )
                 )
-            )
 
             return images, global_labels
             # images.shape = [e*w*(q+s) x c x h x w],  global_labels.shape = [e x w x (q+s)]

core/data/collates/contrib/__init__.py

@@ -67,6 +67,10 @@ def get_augment_method(
                 transforms.RandomHorizontalFlip(),
                 transforms.ColorJitter(**CJ_DICT),
             ]
+        elif config["augment_method"] == "COSOCAugment":
+            trfms_list = [
+                transforms.RandomHorizontalFlip(),
+            ]
         else:
             trfms_list = get_default_image_size_trfms(config["image_size"])
             trfms_list += [
@@ -75,24 +79,30 @@ def get_augment_method(
             ]
 
     else:
-        if config["image_size"] == 224:
-            trfms_list = [
-                transforms.Resize((256, 256)),
-                transforms.CenterCrop((224, 224)),
-            ]
-        elif config["image_size"] == 84:
+        if config['classifier']['name'] == 'COSOC':
             trfms_list = [
-                transforms.Resize((96, 96)),
-                transforms.CenterCrop((84, 84)),
-            ]
-        # for MTL -> alternative solution: use avgpool(ks=11)
-        elif config["image_size"] == 80:
-            trfms_list = [
-                transforms.Resize((92, 92)),
-                transforms.CenterCrop((80, 80)),
+                transforms.RandomResizedCrop(config["image_size"]),
+                transforms.RandomHorizontalFlip(),
             ]
         else:
-            raise RuntimeError
+            if config["image_size"] == 224:
+                trfms_list = [
+                    transforms.Resize((256, 256)),
+                    transforms.CenterCrop((224, 224)),
+                ]
+            elif config["image_size"] == 84:
+                trfms_list = [
+                    transforms.Resize((96, 96)),
+                    transforms.CenterCrop((84, 84)),
+                ]
+            # for MTL -> alternative solution: use avgpool(ks=11)
+            elif config["image_size"] == 80:
+                trfms_list = [
+                    transforms.Resize((92, 92)),
+                    transforms.CenterCrop((80, 80)),
+                ]
+            else:
+                raise RuntimeError
 
     return trfms_list
 

core/data/dataloader.py

@@ -4,7 +4,7 @@
 from torch.utils.data.distributed import DistributedSampler
 from torchvision import transforms
 
-from core.data.dataset import GeneralDataset
+from core.data.dataset import GeneralDataset, COSOCDataset
 from .collates import get_collate_function, get_augment_method,get_mean_std
 from .samplers import DistributedCategoriesSampler, get_sampler
 from ..utils import ModelType
@@ -40,16 +40,27 @@ def get_dataloader(config, mode, model_type, distribute):
     MEAN,STD=get_mean_std(config, mode)
 
     trfms_list = get_augment_method(config, mode)
-
     trfms_list.append(transforms.ToTensor())
     trfms_list.append(transforms.Normalize(mean=MEAN, std=STD))
     trfms = transforms.Compose(trfms_list)
 
-    dataset = GeneralDataset(
-        data_root=config["data_root"],
-        mode=mode,
-        use_memory=config["use_memory"],
-    )
+    if config['classifier']['name'] == 'COSOC':
+        dataset = COSOCDataset(
+            data_root=config["data_root"],
+            mode=mode,
+            use_memory=config["use_memory"],
+            feature_image_and_crop_id=config['feature_image_and_crop_id'],
+            position_list=config['position_list'],
+            # ratio=config['ratio'],
+            # crop_size=config['crop_size'],
+            image_sz=config['image_size'],
+        )
+    else:
+        dataset = GeneralDataset(
+            data_root=config["data_root"],
+            mode=mode,
+            use_memory=config["use_memory"],
+        )
 
     if config["dataloader_num"] == 1 or mode in ["val", "test"]:
 

core/data/dataset.py

@@ -5,6 +5,11 @@
 
 from PIL import Image
 from torch.utils.data import Dataset
+from torchvision import transforms
+import torchvision.transforms.functional as functional
+import numpy as np
+import torch
+import random
 
 
 def pil_loader(path):
@@ -183,3 +188,101 @@ def __getitem__(self, idx):
         label = self.label_list[idx]
 
         return data, label
+
+def crop_func(img, crop, ratio = 1.2):
+    """
+    Given cropping positios, relax for a certain ratio, and return new crops
+    , along with the area ratio.
+    """
+    assert len(crop) == 4
+    w,h = functional.get_image_size(img)
+    if crop[0] == -1.:
+        crop[0],crop[1],crop[2],crop[3]  = 0., 0., h, w
+    else:
+        crop[0] = max(0, crop[0]-crop[2]*(ratio-1)/2)
+        crop[1] = max(0, crop[1]-crop[3]*(ratio-1)/2)
+        crop[2] = min(ratio*crop[2], h-crop[0])
+        crop[3] = min(ratio*crop[3], w-crop[1])
+    return crop, crop[2]*crop[3]/(w*h)
+
+class COSOCDataset(GeneralDataset):
+    def __init__(self, data_root="", mode="train", loader=default_loader, use_memory=True, trfms=None, feature_image_and_crop_id='', position_list='', ratio = 1.2, crop_size = 0.08, image_sz = 84):
+        super().__init__(data_root, mode, loader, use_memory, trfms)
+        self.image_sz = image_sz
+        self.ratio = ratio
+        self.crop_size = crop_size
+        with open(feature_image_and_crop_id, 'rb') as f:
+            self.feature_image_and_crop_id = pickle.load(f)
+        self.position_list = np.load(position_list)
+        self._get_id_position_map()
+
+    def _get_id_position_map(self):
+        self.position_map = {}
+        for i, feature_image_and_crop_ids in self.feature_image_and_crop_id.items():
+            for clusters in feature_image_and_crop_ids:
+                for image in clusters:
+                    # print(image)
+                    if image[0] in self.position_map:
+                        self.position_map[image[0]].append((image[1],image[2]))
+                    else:
+                        self.position_map[image[0]] = [(image[1],image[2])]
+
+    def _multi_crop_get(self, idx):
+        if self.use_memory:
+            data = self.data_list[idx]
+        else:
+            image_name = self.data_list[idx]
+            image_path = os.path.join(self.data_root, "images", image_name)
+            data = self.loader(image_path)
+            ... # image -> aug(collate) -> tensor (b, patch, ...) -> classifier
+
+        if self.trfms is not None:
+            data = self.trfms(data)
+        label = self.label_list[idx]
+
+        return data, label
+
+    def _prob_crop_get(self, idx):
+        if self.use_memory:
+            data = self.data_list[idx]
+        else:
+            image_name = self.data_list[idx]
+            image_path = os.path.join(self.data_root, "images", image_name)
+            data = self.loader(image_path)
+            idx = int(idx)
+
+            x = random.random()
+            ran_crop_prob = 1 - torch.tensor(self.position_map[idx][0][1]).sum()
+            if x > ran_crop_prob:
+                crop_ids = self.position_map[idx][0][0]
+                if ran_crop_prob <= x < ran_crop_prob+self.position_map[idx][0][1][0]:
+                    crop_id = crop_ids[0]
+                elif ran_crop_prob+self.position_map[idx][0][1][0] <= x < ran_crop_prob+self.position_map[idx][0][1][1]+self.position_map[idx][0][1][0]:
+                    crop_id = crop_ids[1]
+                else:
+                    crop_id = crop_ids[2]
+                crop = self.position_list[idx][crop_id]
+                crop, space_ratio = crop_func(data, crop, ratio = self.ratio)
+                data = functional.crop(data,crop[0],crop[1], crop[2],crop[3])
+                data = transforms.RandomResizedCrop(self.image_sz, scale = (self.crop_size/space_ratio, 1.0))(data)
+            else:
+                data = transforms.RandomResizedCrop(self.image_sz)(data)
+
+        if self.trfms is not None:
+            data = self.trfms(data)
+        label = self.label_list[idx]
+        return data, label
+
+    def __getitem__(self, idx):
+        """Return a PyTorch like dataset item of (data, label) tuple.
+
+        Args:
+            idx (int): The __getitem__ id.
+
+        Returns:
+            tuple: A tuple of (image, label)
+        """
+        if self.mode == 'train':
+            return self._prob_crop_get(idx)
+        else:
+            return self._multi_crop_get(idx)

core/model/backbone/__init__.py

@@ -3,6 +3,7 @@
 from .conv_four_mcl import Conv64F_MCL
 from .resnet_12 import resnet12, resnet12woLSC
 from .resnet_12_mcl import resnet12_mcl,resnet12_r2d2
+from .resnet_12_cosoc import resnet12_cosoc
 from .resnet_18 import resnet18
 from .wrn import WRN
 from .resnet_12_mtl_offcial import resnet12MTLofficial
@@ -11,7 +12,6 @@
 from .resnet_bdc import resnet12Bdc, resnet18Bdc
 from core.model.backbone.utils.maml_module import convert_maml_module
 
-
 def get_backbone(config):
     """Get the backbone according to the config dict.
 

core/model/backbone/resnet_12.py

@@ -185,6 +185,7 @@ def __init__(
         maxpool_last2=True,
     ):
         self.inplanes = 3
+        self.outdim = planes[-1]
         super(ResNet, self).__init__()
 
         self.layer1 = self._make_layer(

core/model/backbone/resnet_12_cosoc.py

@@ -0,0 +1,97 @@
+import torch.nn as nn
+
+
+def conv3x3(in_planes, out_planes):
+    return nn.Conv2d(in_planes, out_planes, 3, padding=1, bias=False)
+
+
+def conv1x1(in_planes, out_planes):
+    return nn.Conv2d(in_planes, out_planes, 1, bias=False)
+
+
+def norm_layer(planes):
+    return nn.BatchNorm2d(planes)
+
+
+class Block(nn.Module):
+
+    def __init__(self, inplanes, planes, downsample):
+        super().__init__()
+
+        self.relu = nn.LeakyReLU(0.1)
+
+        self.conv1 = conv3x3(inplanes, planes)
+        self.bn1 = norm_layer(planes)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = norm_layer(planes)
+        self.conv3 = conv3x3(planes, planes)
+        self.bn3 = norm_layer(planes)
+
+        self.downsample = downsample
+
+        self.maxpool = nn.MaxPool2d(2)
+
+    def forward(self, 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)
+
+        identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        out = self.maxpool(out)
+
+        return out
+
+
+class ResNet12(nn.Module):
+    """The standard popular ResNet12 Model used in Few-Shot Learning.
+    """
+    def __init__(self, channels):
+        super().__init__()
+
+        self.inplanes = 3
+
+        self.layer1 = self._make_layer(channels[0])
+        self.layer2 = self._make_layer(channels[1])
+        self.layer3 = self._make_layer(channels[2])
+        self.layer4 = self._make_layer(channels[3])
+
+        self.outdim = channels[3]
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out',
+                                        nonlinearity='leaky_relu')
+            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+    def _make_layer(self, planes):
+        downsample = nn.Sequential(
+            conv1x1(self.inplanes, planes),
+            norm_layer(planes),
+        )
+        block = Block(self.inplanes, planes, downsample)
+        self.inplanes = planes
+        return block
+
+    def forward(self, x):
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        # x = x.view(x.shape[0], x.shape[1], -1).mean(dim=2).unsqueeze_(2).unsqueeze_(3)
+        return x
+
+
+def resnet12_cosoc():
+    return ResNet12([64, 160, 320, 640])

core/model/meta/matchingnet.py

@@ -6,7 +6,6 @@
 from .meta_model import MetaModel
 from core.utils import accuracy
 from ..backbone.utils import convert_maml_module
-import utils
 import torch.nn.functional as F
 
 class IFSLUtils(nn.Module):