Add util, model files and clean vs grean notebook for UC2018

talks/uc2018/Plenary/pools/model.py

@@ -0,0 +1,211 @@
+import numpy as np
+from PIL import Image
+
+import torch
+from torch import nn,optim
+from torchvision.models import *
+from torch.autograd import Variable
+import torch.nn.functional as F
+from torch.nn.init import kaiming_normal
+import os
+
+import json, pdb
+from PIL import ImageDraw, ImageFont
+from matplotlib import patches, patheffects
+# torch.cuda.set_device(0)
+from distutils.version import LooseVersion
+
+import matplotlib.cm as cmx
+import matplotlib.colors as mcolors
+from cycler import cycler
+from PIL import ImageOps
+import torch.nn.functional as F
+
+total_classes = 2
+k = 9
+
+def children(m): return m if isinstance(m, (list, tuple)) else list(m.children())
+
+def apply_init(m, init_fn):
+    m.apply(lambda x: cond_init(x, init_fn))
+
+def cond_init(m, init_fn):
+    if not isinstance(m, (nn.BatchNorm1d,nn.BatchNorm2d,nn.BatchNorm3d)):
+        if hasattr(m, 'weight'): init_fn(m.weight)
+    if hasattr(m, 'bias') and hasattr(m.bias, 'data'): m.bias.data.fill_(0.)
+
+class StdConv(nn.Module):
+    def __init__(self, nin, nout, stride=2, drop=0.1):
+        super().__init__()
+        self.conv = nn.Conv2d(nin, nout, 3, stride=stride, padding=1)
+        self.bn = nn.BatchNorm2d(nout)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x): return self.drop(self.bn(F.relu(self.conv(x))))
+
+def flatten_conv(x,k):
+    bs,nf,gx,gy = x.size()                # batch size, num filters, width, height
+    x = x.permute(0,3,2,1).contiguous()
+    return x.view(bs,-1,nf//k)
+
+class OutConv(nn.Module):
+    def __init__(self, k, nin, bias):
+        super().__init__()
+        self.k = k
+        self.oconv1 = nn.Conv2d(nin, (total_classes + 1)*k, 3, padding=1) # nclasses
+        self.oconv2 = nn.Conv2d(nin, 4*k, 3, padding=1) # bboxes
+        self.oconv1.bias.data.zero_().add_(bias)
+
+    def forward(self, x):
+        return [flatten_conv(self.oconv1(x), self.k),
+                flatten_conv(self.oconv2(x), self.k)]
+
+drop=0.4
+
+class SSD_MultiHead(nn.Module):
+    def __init__(self, k, bias):
+        super().__init__()
+        self.drop = nn.Dropout(drop)
+        self.sconv0 = StdConv(512,256, stride=1, drop=drop)
+        self.sconv1 = StdConv(256,256, drop=drop)
+        self.sconv2 = StdConv(256,256, drop=drop)
+        self.sconv3 = StdConv(256,256, drop=drop)
+        self.out0 = OutConv(k, 256, bias)
+        self.out1 = OutConv(k, 256, bias)
+        self.out2 = OutConv(k, 256, bias)
+        self.out3 = OutConv(k, 256, bias)
+
+    def forward(self, x):
+        x = self.drop(F.relu(x))
+        x = self.sconv0(x)
+        x = self.sconv1(x)
+        o1c,o1l = self.out1(x)
+        x = self.sconv2(x)
+        o2c,o2l = self.out2(x)
+        x = self.sconv3(x)
+        o3c,o3l = self.out3(x)
+        return [torch.cat([o1c,o2c,o3c], dim=1),
+                torch.cat([o1l,o2l,o3l], dim=1)]
+
+
+
+IS_TORCH_04 = LooseVersion(torch.__version__) >= LooseVersion('0.4')
+
+model_meta = {
+    resnet18:[8,6], resnet34:[8,6], resnet50:[8,6], resnet101:[8,6], resnet152:[8,6],
+    vgg16:[0,22], vgg19:[0,22]
+}
+
+requires_grad = False
+USE_GPU = True
+
+def map_over(x, f): return [f(o) for o in x] if is_listy(x) else f(x)
+def V (x, requires_grad=False, volatile=False): return map_over(x, lambda o: V_(o, requires_grad, volatile))
+def V_(x, requires_grad=False, volatile=False): return create_variable(x, volatile, requires_grad)
+def is_listy(x): return isinstance(x, (list,tuple))
+def A(*a): return np.array(a[0]) if len(a)==1 else [np.array(o) for o in a]
+
+def create_variable(x, volatile, requires_grad=False):
+    if type (x) != Variable:
+        if IS_TORCH_04: x = Variable(T(x), requires_grad=requires_grad)
+        else:           x = Variable(T(x), requires_grad=requires_grad, volatile=volatile)
+    return x
+
+def T(a, half=False, cuda=True):
+    if not torch.is_tensor(a):
+        a = np.array(np.ascontiguousarray(a))
+        if a.dtype in (np.int8, np.int16, np.int32, np.int64):
+            a = torch.LongTensor(a.astype(np.int64))
+        elif a.dtype in (np.float32, np.float64):
+            a = torch.cuda.HalfTensor(a) if half else torch.FloatTensor(a)
+        else: raise NotImplementedError(a.dtype)
+    if cuda: a = to_gpu(a, async=True)
+    return a
+
+def to_gpu(x, *args, **kwargs):
+    return x.cuda(*args, **kwargs) if USE_GPU else x
+
+def to_np(v):
+    if isinstance(v, (np.ndarray, np.generic)): return v
+    if isinstance(v, (list,tuple)): return [to_np(o) for o in v]
+    if isinstance(v, Variable): v=v.data
+    if isinstance(v, torch.cuda.HalfTensor): v=v.float()
+    return v.cpu().numpy()
+
+def cut_model(m, cut):
+    return list(m.children())[:cut] if cut else [m]
+
+class AdaptiveConcatPool2d(nn.Module):
+    def __init__(self, sz=None):
+        super().__init__()
+        sz = sz or (1,1)
+        self.ap = nn.AdaptiveAvgPool2d(sz)
+        self.mp = nn.AdaptiveMaxPool2d(sz)
+    def forward(self, x): return torch.cat([self.mp(x), self.ap(x)], 1)
+
+class Flatten(nn.Module):
+    def __init__(self): super().__init__()
+    def forward(self, x): return x.view(x.size(0), -1)
+
+def num_features(m):
+    c=children(m)
+    if len(c)==0: return None
+    for l in reversed(c):
+        if hasattr(l, 'num_features'): return l.num_features
+        res = num_features(l)
+        if res is not None: return res  
+
+class ConvnetBuilder():
+
+    def __init__(self, f, c, is_multi, is_reg, ps=None, xtra_fc=None, xtra_cut=0, custom_head=None, pretrained=True):
+        self.f,self.c,self.is_multi,self.is_reg,self.xtra_cut = f,c,is_multi,is_reg,xtra_cut
+        if xtra_fc is None: xtra_fc = [512]
+        if ps is None: ps = [0.25]*len(xtra_fc) + [0.5]
+        self.ps,self.xtra_fc = ps,xtra_fc
+
+        if f in model_meta: cut,self.lr_cut = model_meta[f]
+        else: cut,self.lr_cut = 0,0
+        cut-=xtra_cut
+        layers = cut_model(f(pretrained), cut)
+        self.nf = num_features(layers)*2
+        if not custom_head: layers += [AdaptiveConcatPool2d(), Flatten()]
+        self.top_model = nn.Sequential(*layers)
+        n_fc = len(self.xtra_fc)+1
+        if not isinstance(self.ps, list): self.ps = [self.ps]*n_fc
+
+        if custom_head: fc_layers = [custom_head]
+        else: fc_layers = self.get_fc_layers()
+        self.n_fc = len(fc_layers)
+        self.fc_model = to_gpu(nn.Sequential(*fc_layers))
+        if not custom_head: apply_init(self.fc_model, kaiming_normal)
+        self.model = to_gpu(nn.Sequential(*(layers+fc_layers)))
+
+    @property
+    def name(self): return f'{self.f.__name__}_{self.xtra_cut}'
+
+    def create_fc_layer(self, ni, nf, p, actn=None):
+        res=[nn.BatchNorm1d(num_features=ni)]
+        if p: res.append(nn.Dropout(p=p))
+        res.append(nn.Linear(in_features=ni, out_features=nf))
+        if actn: res.append(actn)
+        return res
+
+    def get_fc_layers(self):
+        res=[]
+        ni=self.nf
+        for i,nf in enumerate(self.xtra_fc):
+            res += self.create_fc_layer(ni, nf, p=self.ps[i], actn=nn.ReLU())
+            ni=nf
+        final_actn = nn.Sigmoid() if self.is_multi else nn.LogSoftmax(1)
+        if self.is_reg: final_actn = None
+        res += self.create_fc_layer(ni, self.c, p=self.ps[-1], actn=final_actn)
+        return res
+
+    def get_layer_groups(self, do_fc=False):
+        if do_fc:
+            return [self.fc_model]
+        idxs = [self.lr_cut]
+        c = children(self.top_model)
+        if len(c)==3: c = children(c[0])+c[1:]
+        lgs = list(split_by_idxs(c,idxs))
+        return lgs+[self.fc_model]