From 071a59d9bcfabf4acc82cbe5f4906ffa4bcca501 Mon Sep 17 00:00:00 2001 From: Krzysztof Lis Date: Sat, 20 Jan 2018 14:25:50 +0100 Subject: [PATCH 1/4] Use automatic downloads of pretrained weights PyTorch will automatically download pretrained weights into `os.environ['TORCH_MODEL_ZOO']` using the mechanism described here: (http://pytorch.org/docs/master/model_zoo.html) Removed hardcoded paths --- models/config.py | 14 ++++---------- models/duc_hdc.py | 7 +------ models/fcn16s.py | 4 +--- models/fcn32s.py | 5 +---- models/fcn8s.py | 18 ++++++++++-------- models/gcn.py | 5 +---- models/psp_net.py | 5 +---- models/seg_net.py | 5 +---- 8 files changed, 20 insertions(+), 43 deletions(-) diff --git a/models/config.py b/models/config.py index 6b27bec..f01b032 100644 --- a/models/config.py +++ b/models/config.py @@ -1,17 +1,11 @@ import os -# here (https://github.com/pytorch/vision/tree/master/torchvision/models) to find the download link of pretrained models - -root = '/media/b3-542/LIBRARY/ZijunDeng/PyTorch Pretrained' -res101_path = os.path.join(root, 'ResNet', 'resnet101-5d3b4d8f.pth') -res152_path = os.path.join(root, 'ResNet', 'resnet152-b121ed2d.pth') -inception_v3_path = os.path.join(root, 'Inception', 'inception_v3_google-1a9a5a14.pth') -vgg19_bn_path = os.path.join(root, 'VggNet', 'vgg19_bn-c79401a0.pth') -vgg16_path = os.path.join(root, 'VggNet', 'vgg16-397923af.pth') -dense201_path = os.path.join(root, 'DenseNet', 'densenet201-4c113574.pth') +# PyTorch will automatically download pretrained weights into `os.environ['TORCH_MODEL_ZOO']` +# using the mechanism described here: (http://pytorch.org/docs/master/model_zoo.html) +# Download links used are also listed here: (https://github.com/pytorch/vision/tree/master/torchvision/models) ''' vgg16 trained using caffe visit this (https://github.com/jcjohnson/pytorch-vgg) to download the converted vgg16 ''' -vgg16_caffe_path = os.path.join(root, 'VggNet', 'vgg16-caffe.pth') +vgg16_caffe_path = os.path.join(os.environ.get('TORCH_MODEL_ZOO', '.'), 'vgg16-caffe.pth') diff --git a/models/duc_hdc.py b/models/duc_hdc.py index 23bf8a4..19a8830 100644 --- a/models/duc_hdc.py +++ b/models/duc_hdc.py @@ -2,9 +2,6 @@ from torch import nn from torchvision import models -from .config import res152_path - - class _DenseUpsamplingConvModule(nn.Module): def __init__(self, down_factor, in_dim, num_classes): super(_DenseUpsamplingConvModule, self).__init__() @@ -26,9 +23,7 @@ class ResNetDUC(nn.Module): # the size of image should be multiple of 8 def __init__(self, num_classes, pretrained=True): super(ResNetDUC, self).__init__() - resnet = models.resnet152() - if pretrained: - resnet.load_state_dict(torch.load(res152_path)) + resnet = models.resnet152(pretrained=pretrained) self.layer0 = nn.Sequential(resnet.conv1, resnet.bn1, resnet.relu, resnet.maxpool) self.layer1 = resnet.layer1 self.layer2 = resnet.layer2 diff --git a/models/fcn16s.py b/models/fcn16s.py index f9873ae..4030442 100644 --- a/models/fcn16s.py +++ b/models/fcn16s.py @@ -9,9 +9,7 @@ class FCN16VGG(nn.Module): def __init__(self, num_classes, pretrained=True): super(FCN16VGG, self).__init__() - vgg = models.vgg16() - if pretrained: - vgg.load_state_dict(torch.load(vgg16_caffe_path)) + vgg = models.vgg16(pretrained=pretrained) features, classifier = list(vgg.features.children()), list(vgg.classifier.children()) features[0].padding = (100, 100) diff --git a/models/fcn32s.py b/models/fcn32s.py index f9e1eb0..f20f9c9 100644 --- a/models/fcn32s.py +++ b/models/fcn32s.py @@ -3,15 +3,12 @@ from torchvision import models from ..utils import get_upsampling_weight -from .config import vgg16_caffe_path class FCN32VGG(nn.Module): def __init__(self, num_classes, pretrained=True): super(FCN32VGG, self).__init__() - vgg = models.vgg16() - if pretrained: - vgg.load_state_dict(torch.load(vgg16_caffe_path)) + vgg = models.vgg16(pretrained=pretrained) features, classifier = list(vgg.features.children()), list(vgg.classifier.children()) features[0].padding = (100, 100) diff --git a/models/fcn8s.py b/models/fcn8s.py index 8492868..165f442 100644 --- a/models/fcn8s.py +++ b/models/fcn8s.py @@ -3,20 +3,22 @@ from torchvision import models from ..utils import get_upsampling_weight -from .config import vgg16_path, vgg16_caffe_path +from .config import vgg16_caffe_path # This is implemented in full accordance with the original one (https://github.com/shelhamer/fcn.berkeleyvision.org) class FCN8s(nn.Module): def __init__(self, num_classes, pretrained=True, caffe=False): super(FCN8s, self).__init__() - vgg = models.vgg16() - if pretrained: - if caffe: - # load the pretrained vgg16 used by the paper's author - vgg.load_state_dict(torch.load(vgg16_caffe_path)) - else: - vgg.load_state_dict(torch.load(vgg16_path)) + + if pretrained and caffe: + vgg = models.vgg16() + # load the pretrained vgg16 used by the paper's author + vgg.load_state_dict(torch.load(vgg16_caffe_path)) + else: + # if pretrained, load the weights from PyTorch model zoo + vgg = models.vgg16(pretrained=pretrained) + features, classifier = list(vgg.features.children()), list(vgg.classifier.children()) ''' diff --git a/models/gcn.py b/models/gcn.py index bc0d6a1..27d5a60 100644 --- a/models/gcn.py +++ b/models/gcn.py @@ -4,7 +4,6 @@ from torchvision import models from ..utils import initialize_weights -from .config import res152_path # many are borrowed from https://github.com/ycszen/pytorch-ss/blob/master/gcn.py @@ -52,9 +51,7 @@ class GCN(nn.Module): def __init__(self, num_classes, input_size, pretrained=True): super(GCN, self).__init__() self.input_size = input_size - resnet = models.resnet152() - if pretrained: - resnet.load_state_dict(torch.load(res152_path)) + resnet = models.resnet152(pretrained=pretrained) self.layer0 = nn.Sequential(resnet.conv1, resnet.bn1, resnet.relu) self.layer1 = nn.Sequential(resnet.maxpool, resnet.layer1) self.layer2 = resnet.layer2 diff --git a/models/psp_net.py b/models/psp_net.py index 7dedea5..577714d 100644 --- a/models/psp_net.py +++ b/models/psp_net.py @@ -5,7 +5,6 @@ from ..utils import initialize_weights from ..utils.misc import Conv2dDeformable -from .config import res101_path class _PyramidPoolingModule(nn.Module): @@ -34,9 +33,7 @@ class PSPNet(nn.Module): def __init__(self, num_classes, pretrained=True, use_aux=True): super(PSPNet, self).__init__() self.use_aux = use_aux - resnet = models.resnet101() - if pretrained: - resnet.load_state_dict(torch.load(res101_path)) + resnet = models.resnet101(pretrained=pretrained) self.layer0 = nn.Sequential(resnet.conv1, resnet.bn1, resnet.relu, resnet.maxpool) self.layer1, self.layer2, self.layer3, self.layer4 = resnet.layer1, resnet.layer2, resnet.layer3, resnet.layer4 diff --git a/models/seg_net.py b/models/seg_net.py index 61128a8..403f071 100644 --- a/models/seg_net.py +++ b/models/seg_net.py @@ -3,7 +3,6 @@ from torchvision import models from ..utils import initialize_weights -from .config import vgg19_bn_path class _DecoderBlock(nn.Module): @@ -35,9 +34,7 @@ def forward(self, x): class SegNet(nn.Module): def __init__(self, num_classes, pretrained=True): super(SegNet, self).__init__() - vgg = models.vgg19_bn() - if pretrained: - vgg.load_state_dict(torch.load(vgg19_bn_path)) + vgg = models.vgg19_bn(pretrained=pretrained) features = list(vgg.features.children()) self.enc1 = nn.Sequential(*features[0:7]) self.enc2 = nn.Sequential(*features[7:14]) From 1c17c23bc8424e7e583cc86b8c7105240c58fe69 Mon Sep 17 00:00:00 2001 From: Krzysztof Maciej Lis Date: Fri, 16 Nov 2018 15:27:23 +0100 Subject: [PATCH 2/4] bayesian segnet --- .gitignore | 2 + models/__init__.py | 1 + models/psp_net.py | 13 +++--- models/seg_net.py | 2 +- models/seg_net_bayes.py | 99 +++++++++++++++++++++++++++++++++++++++++ utils/misc.py | 2 +- 6 files changed, 111 insertions(+), 8 deletions(-) create mode 100644 .gitignore create mode 100644 models/seg_net_bayes.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..badb8ed --- /dev/null +++ b/.gitignore @@ -0,0 +1,2 @@ + __pycache__/ + models/seg_net_bayes.py diff --git a/models/__init__.py b/models/__init__.py index 0191f88..1663862 100644 --- a/models/__init__.py +++ b/models/__init__.py @@ -5,4 +5,5 @@ from .gcn import * from .psp_net import * from .seg_net import * +from .seg_net_bayes import * from .u_net import * diff --git a/models/psp_net.py b/models/psp_net.py index 577714d..154d201 100644 --- a/models/psp_net.py +++ b/models/psp_net.py @@ -9,7 +9,7 @@ class _PyramidPoolingModule(nn.Module): def __init__(self, in_dim, reduction_dim, setting): - super(_PyramidPoolingModule, self).__init__() + super().__init__() self.features = [] for s in setting: self.features.append(nn.Sequential( @@ -24,14 +24,14 @@ def forward(self, x): x_size = x.size() out = [x] for f in self.features: - out.append(F.upsample(f(x), x_size[2:], mode='bilinear')) + out.append(F.interpolate(f(x), x_size[2:], mode='bilinear')) out = torch.cat(out, 1) return out class PSPNet(nn.Module): def __init__(self, num_classes, pretrained=True, use_aux=True): - super(PSPNet, self).__init__() + super().__init__() self.use_aux = use_aux resnet = models.resnet101(pretrained=pretrained) self.layer0 = nn.Sequential(resnet.conv1, resnet.bn1, resnet.relu, resnet.maxpool) @@ -63,7 +63,8 @@ def __init__(self, num_classes, pretrained=True, use_aux=True): initialize_weights(self.ppm, self.final) - def forward(self, x): + def forward(self, image): + x = image x_size = x.size() x = self.layer0(x) x = self.layer1(x) @@ -75,8 +76,8 @@ def forward(self, x): x = self.ppm(x) x = self.final(x) if self.training and self.use_aux: - return F.upsample(x, x_size[2:], mode='bilinear'), F.upsample(aux, x_size[2:], mode='bilinear') - return F.upsample(x, x_size[2:], mode='bilinear') + return F.interpolate(x, x_size[2:], mode='bilinear'), F.interpolate(aux, x_size[2:], mode='bilinear') + return F.interpolate(x, x_size[2:], mode='bilinear') # just a try, not recommend to use diff --git a/models/seg_net.py b/models/seg_net.py index 403f071..655833b 100644 --- a/models/seg_net.py +++ b/models/seg_net.py @@ -8,7 +8,7 @@ class _DecoderBlock(nn.Module): def __init__(self, in_channels, out_channels, num_conv_layers): super(_DecoderBlock, self).__init__() - middle_channels = in_channels / 2 + middle_channels = in_channels // 2 layers = [ nn.ConvTranspose2d(in_channels, in_channels, kernel_size=2, stride=2), nn.Conv2d(in_channels, middle_channels, kernel_size=3, padding=1), diff --git a/models/seg_net_bayes.py b/models/seg_net_bayes.py new file mode 100644 index 0000000..a7f8f04 --- /dev/null +++ b/models/seg_net_bayes.py @@ -0,0 +1,99 @@ + +import torch +from torch import nn +from torchvision import models + +from ..utils import initialize_weights +from .seg_net import _DecoderBlock, SegNet + +class SegNetBayes(SegNet): + def __init__(self, num_classes, dropout_p=0.5, pretrained=True, num_samples=16, min_batch_size=4): + """ + @param num_samples: number of samples for the Monte-Carlo simulation, + how many times to run with random dropout + """ + super().__init__(num_classes=num_classes, pretrained=pretrained) + + self.drop = nn.Dropout2d(p=dropout_p, inplace=False) + self.num_samples = num_samples + self.min_batch_size = min_batch_size + + def forward(self, x): + enc1 = self.enc1(x) + #print('enc1', enc1.shape) + + enc2 = self.enc2(enc1) + #print('enc2', enc2.shape) + + enc3 = self.enc3(enc2) + #print('enc3', enc3.shape) + enc3 = self.drop(enc3) + #print('enc3d', enc3.shape) + + enc4 = self.enc4(enc3) + #print('enc4', enc4.shape) + enc4 = self.drop(enc4) + #print('enc4d', enc4.shape) + + enc5 = self.enc5(enc4) + #print('enc5', enc5.shape) + enc5 = self.drop(enc5) + #print('enc5d', enc5.shape) + + dec5 = self.dec5(enc5) + #print('dec5', dec5.shape) + dec5 = self.drop(dec5) + #print('dec5d', dec5.shape) + + dec4 = self.dec4(torch.cat([enc4, dec5], 1)) + #print('dec4', dec4.shape) + dec4 = self.drop(dec4) + #print('dec4d', dec4.shape) + + dec3 = self.dec3(torch.cat([enc3, dec4], 1)) + dec3 = self.drop(dec3) + + dec2 = self.dec2(torch.cat([enc2, dec3], 1)) + dec1 = self.dec1(torch.cat([enc1, dec2], 1)) + return dec1 + + def forward_multisample(self, x, num_samples=None): + # dropout must be active + backup_train_mode = self.drop.training + self.drop.train() + + softmax = torch.nn.Softmax2d() + + num_samples = num_samples if num_samples else self.num_samples + + results = [softmax(self.forward(x)).data.cpu() for i in range(num_samples)] + + preds = torch.stack(results).cuda() + avg = torch.mean(preds, 0) + var = torch.var(preds, 0) + del preds + + # restore mode + self.drop.train(backup_train_mode) + + return dict( + mean = avg, + var = var, + ) + + #def sample(self, x, num_samples, batch_size): + #infer desired batch size from input shape + #we will divide a num_samples into batches + #num_frames = x.shape[0] + #batch_size = max(num_frames, self.min_batch_size) + + #for sample_idx in range(num_samples): + #pred = + + + #for fr_idx in range(num_frames): + #x_single = x[fr_idx:fr_idx+1, :, :, :] + #self.sample(x_single, num_samples, batch_size) + + + diff --git a/utils/misc.py b/utils/misc.py index 90e3c19..06aa33d 100644 --- a/utils/misc.py +++ b/utils/misc.py @@ -17,7 +17,7 @@ def initialize_weights(*models): for model in models: for module in model.modules(): if isinstance(module, nn.Conv2d) or isinstance(module, nn.Linear): - nn.init.kaiming_normal(module.weight) + nn.init.kaiming_normal_(module.weight) if module.bias is not None: module.bias.data.zero_() elif isinstance(module, nn.BatchNorm2d): From f702875fe9f4c197c9fc8744ab94c542f90d0abe Mon Sep 17 00:00:00 2001 From: Krzysztof Maciej Lis Date: Fri, 25 Oct 2019 09:54:24 +0200 Subject: [PATCH 3/4] psp multihead --- .gitignore | 6 ++- models/__init__.py | 1 + models/psp_net_multihead.py | 75 +++++++++++++++++++++++++++++++++++++ 3 files changed, 80 insertions(+), 2 deletions(-) create mode 100644 models/psp_net_multihead.py diff --git a/.gitignore b/.gitignore index badb8ed..581f124 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,4 @@ - __pycache__/ - models/seg_net_bayes.py + +__pycache__/ +models/seg_net_bayes.py + diff --git a/models/__init__.py b/models/__init__.py index 1663862..2e00774 100644 --- a/models/__init__.py +++ b/models/__init__.py @@ -4,6 +4,7 @@ from .fcn8s import * from .gcn import * from .psp_net import * +from .psp_net_multihead import * from .seg_net import * from .seg_net_bayes import * from .u_net import * diff --git a/models/psp_net_multihead.py b/models/psp_net_multihead.py new file mode 100644 index 0000000..d380ecf --- /dev/null +++ b/models/psp_net_multihead.py @@ -0,0 +1,75 @@ +import torch +import torch.nn.functional as F +from torch import nn +from torchvision import models + +from ..utils import initialize_weights + +from .psp_net import _PyramidPoolingModule + + +class PSPHead(nn.Module): + def __init__(self, num_classes): + super().__init__() + + self.ppm = _PyramidPoolingModule(2048, 512, (1, 2, 3, 6)) + self.final = nn.Sequential( + nn.Conv2d(4096, 512, kernel_size=3, padding=1, bias=False), + nn.BatchNorm2d(512, momentum=.95), + nn.ReLU(inplace=True), + nn.Dropout(0.1), + nn.Conv2d(512, num_classes, kernel_size=1) + ) + + initialize_weights(self.ppm, self.final) + + def forward(self, features_from_backbone, img_size): + result = self.final(self.ppm(features_from_backbone)) + return F.interpolate(result, img_size[2:], mode='bilinear') + + +class PSPNet_Multihead(nn.Module): + def __init__(self, num_heads, num_classes, pretrained=True): + super().__init__() + + self.init_heads(num_heads, num_classes=num_classes) + self.init_backbone(pretrained=pretrained) + + def init_backbone(self, pretrained): + resnet = models.resnet101(pretrained=pretrained) + + for n, m in resnet.layer3.named_modules(): + if 'conv2' in n: + m.dilation, m.padding, m.stride = (2, 2), (2, 2), (1, 1) + elif 'downsample.0' in n: + m.stride = (1, 1) + for n, m in resnet.layer4.named_modules(): + if 'conv2' in n: + m.dilation, m.padding, m.stride = (4, 4), (4, 4), (1, 1) + elif 'downsample.0' in n: + m.stride = (1, 1) + + self.backbone = nn.Sequential( + nn.Sequential(resnet.conv1, resnet.bn1, resnet.relu, resnet.maxpool), # layer 0 + resnet.layer1, + resnet.layer2, + resnet.layer3, + resnet.layer4, + ) + + def init_heads(self, num_heads, num_classes): + + self.heads = nn.Sequential( + *[PSPHead(num_classes=num_classes) for i in range(num_heads)] + ) + + def forward(self, image): + img_size = image.size() + + backbone_features = self.backbone(image) + + return torch.cat([ + head(backbone_features, img_size=img_size) + for head in self.heads + ], dim=1) + From 44d1784984cfd0926821c3fdbc20d371bb074296 Mon Sep 17 00:00:00 2001 From: Krzysztof Lis Date: Fri, 25 Oct 2019 11:05:27 +0200 Subject: [PATCH 4/4] start without skimage --- utils/transforms.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/utils/transforms.py b/utils/transforms.py index 65fcbae..52044f6 100644 --- a/utils/transforms.py +++ b/utils/transforms.py @@ -1,7 +1,10 @@ import random import numpy as np -from skimage.filters import gaussian +try: + from skimage.filters import gaussian +except: + pass import torch from PIL import Image, ImageFilter