def densenet161(pretrained=False, **kwargs):
r"""Densenet-161 model from
`"Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993.pdf>`_
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
"""
model = DenseNet(num_init_features=96, growth_rate=48, block_config=(6, 12, 36, 24),
**kwargs)
if pretrained:
# '.'s are no longer allowed in module names, but pervious _DenseLayer
# has keys 'norm.1', 'relu.1', 'conv.1', 'norm.2', 'relu.2', 'conv.2'.
# They are also in the checkpoints in model_urls. This pattern is used
# to find such keys.
pattern = re.compile(
r'^(.*denselayer\d+\.(?:norm|relu|conv))\.((?:[12])\.(?:weight|bias|running_mean|running_var))$')
state_dict = model_zoo.load_url(model_urls['densenet161'])
for key in list(state_dict.keys()):
res = pattern.match(key)
if res:
new_key = res.group(1) + res.group(2)
state_dict[new_key] = state_dict[key]
del state_dict[key]
model.load_state_dict(state_dict)
return model
def resnet(name, **kwargs):
pretrained_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'
}
blocks = {
'resnet18': BasicBlock,
'resnet34': BasicBlock,
'resnet50': Bottleneck,
'resnet101': Bottleneck,
'resnet152': Bottleneck
}
layers = {
'resnet18': [2, 2, 2, 2],
'resnet34': [3, 4, 6, 3],
'resnet50': [3, 4, 6, 3],
'resnet101': [3, 4, 23, 3],
'resnet152': [3, 8, 36, 3]
}
model = ResNet(blocks[name], layers[name], **kwargs).cuda()
sc.convert(model, model_zoo.load_url(pretrained_urls[name]))
return model
def load_state_dict(model, model_urls, model_root):
from torch.utils import model_zoo
from torch import nn
import re
from collections import OrderedDict
own_state_old = model.state_dict()
own_state = OrderedDict() # remove all 'group' string
for k, v in own_state_old.items():
k = re.sub('group\d+\.', '', k)
own_state[k] = v
state_dict = model_zoo.load_url(model_urls, model_root)
for name, param in state_dict.items():
if name not in own_state:
print(own_state.keys())
raise KeyError('unexpected key "{}" in state_dict'
.format(name))
if isinstance(param, nn.Parameter):
# backwards compatibility for serialized parameters
param = param.data
own_state[name].copy_(param)
missing = set(own_state.keys()) - set(state_dict.keys())
if len(missing) > 0:
raise KeyError('missing keys in state_dict: "{}"'.format(missing))
def inceptionresnetv2(num_classes=1000, pretrained='imagenet'):
r"""InceptionResNetV2 model architecture from the
`"InceptionV4, Inception-ResNet..." <https://arxiv.org/abs/1602.07261>`_ paper.
"""
if pretrained:
settings = pretrained_settings['inceptionresnetv2'][pretrained]
assert num_classes == settings['num_classes'], \
"num_classes should be {}, but is {}".format(settings['num_classes'], num_classes)
# both 'imagenet'&'imagenet+background' are loaded from same parameters
model = InceptionResNetV2(num_classes=1001)
model.load_state_dict(model_zoo.load_url(settings['url']))
if pretrained == 'imagenet':
new_last_linear = nn.Linear(1536, 1000)
new_last_linear.weight.data = model.last_linear.weight.data[1:]
new_last_linear.bias.data = model.last_linear.bias.data[1:]
model.last_linear = new_last_linear
model.input_space = settings['input_space']
model.input_size = settings['input_size']
model.input_range = settings['input_range']
model.mean = settings['mean']
model.std = settings['std']
else:
model = InceptionResNetV2(num_classes=num_classes)
return model
def nasnetalarge(num_classes=1000, pretrained='imagenet'):
r"""NASNetALarge model architecture from the
`"NASNet" <https://arxiv.org/abs/1707.07012>`_ paper.
"""
if pretrained:
settings = pretrained_settings['nasnetalarge'][pretrained]
assert num_classes == settings['num_classes'], \
"num_classes should be {}, but is {}".format(settings['num_classes'], num_classes)
# both 'imagenet'&'imagenet+background' are loaded from same parameters
model = NASNetALarge(num_classes=1001)
model.load_state_dict(model_zoo.load_url(settings['url']))
if pretrained == 'imagenet':
new_last_linear = nn.Linear(model.last_linear.in_features, 1000)
new_last_linear.weight.data = model.last_linear.weight.data[1:]
new_last_linear.bias.data = model.last_linear.bias.data[1:]
model.last_linear = new_last_linear
model.input_space = settings['input_space']
model.input_size = settings['input_size']
model.input_range = settings['input_range']
model.mean = settings['mean']
model.std = settings['std']
else:
model = NASNetALarge(num_classes=num_classes)
return model
def _load_xception_pretrained(self):
pretrain_dict = model_zoo.load_url('http://data.lip6.fr/cadene/pretrainedmodels/xception-b5690688.pth')
model_dict = {}
state_dict = self.state_dict()
for k, v in pretrain_dict.items():
if k in state_dict:
if 'pointwise' in k:
v = v.unsqueeze(-1).unsqueeze(-1)
if k.startswith('block12'):
model_dict[k.replace('block12', 'block20')] = v
elif k.startswith('block11'):
model_dict[k.replace('block11', 'block12')] = v
model_dict[k.replace('block11', 'block13')] = v
model_dict[k.replace('block11', 'block14')] = v
model_dict[k.replace('block11', 'block15')] = v
model_dict[k.replace('block11', 'block16')] = v
model_dict[k.replace('block11', 'block17')] = v
model_dict[k.replace('block11', 'block18')] = v
model_dict[k.replace('block11', 'block19')] = v
elif k.startswith('conv3'):
model_dict[k] = v
elif k.startswith('bn3'):
model_dict[k] = v
model_dict[k.replace('bn3', 'bn4')] = v
elif k.startswith('conv4'):
model_dict[k.replace('conv4', 'conv5')] = v
elif k.startswith('bn4'):
model_dict[k.replace('bn4', 'bn5')] = v
else:
model_dict[k] = v
state_dict.update(model_dict)
self.load_state_dict(state_dict)
def flow_resnet50_aux(pretrained=False, **kwargs):
"""Constructs a ResNet-50 model.
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
"""
model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
if pretrained:
# model.load_state_dict(model_zoo.load_url(model_urls['resnet50']))
pretrained_dict = model_zoo.load_url(model_urls['resnet50'])
model_dict = model.state_dict()
fc_origin_weight = pretrained_dict["fc.weight"].data.numpy()
fc_origin_bias = pretrained_dict["fc.bias"].data.numpy()
# 1. filter out unnecessary keys
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
# 2. overwrite entries in the existing state dict
model_dict.update(pretrained_dict)
# print(model_dict)
fc_new_weight = model_dict["fc_aux.weight"].numpy()
fc_new_bias = model_dict["fc_aux.bias"].numpy()
fc_new_weight[:1000, :] = fc_origin_weight
fc_new_bias[:1000] = fc_origin_bias
model_dict["fc_aux.weight"] = torch.from_numpy(fc_new_weight)
model_dict["fc_aux.bias"] = torch.from_numpy(fc_new_bias)
# 3. load the new state dict
model.load_state_dict(model_dict)
return model
def test_super_resolution(self):
super_resolution_net = SuperResolutionNet(upscale_factor=3)
state_dict = model_zoo.load_url(model_urls['super_resolution'], progress=False)
x = Variable(torch.randn(1, 1, 224, 224), requires_grad=True)
self.run_model_test(super_resolution_net, train=False,
batch_size=BATCH_SIZE, state_dict=state_dict,
input=x, use_gpu=False, atol=1e-6)
def test_srresnet(self):
super_resolution_net = SRResNet(
rescale_factor=4, n_filters=64, n_blocks=8)
state_dict = model_zoo.load_url(model_urls['srresNet'], progress=False)
x = Variable(torch.randn(1, 3, 224, 224), requires_grad=True)
self.run_model_test(super_resolution_net, train=False,
batch_size=1, state_dict=state_dict,
input=x, use_gpu=False)
def init_params(self):
"""Load ImageNet pretrained weights"""
settings = pretrained_settings['nasnetamobile']['imagenet']
pretrained_dict = model_zoo.load_url(settings['url'], map_location=None)
model_dict = self.state_dict()
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
model_dict.update(pretrained_dict)
self.load_state_dict(model_dict)
def create_mtcnn_net(self, use_cuda=True):
self.device = torch.device(
"cuda" if use_cuda and torch.cuda.is_available() else "cpu")
pnet = PNet()
pnet.load_state_dict(model_zoo.load_url(model_urls['pnet']))
pnet.to(self.device).eval()
onet = ONet()
onet.load_state_dict(model_zoo.load_url(model_urls['onet']))
onet.to(self.device).eval()
rnet = RNet()
rnet.load_state_dict(model_zoo.load_url(model_urls['rnet']))
rnet.to(self.device).eval()
return pnet, rnet, onet
def ResNet18_imagenet(pretrained=False, **kwargs):
"""Constructs a ResNet-18 model.
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
"""
model = ResNet_imagenet(BasicBlock, [2, 2, 2, 2], **kwargs)
if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['resnet18']))
return model
def ResNet101_imagenet(pretrained=False, **kwargs):
"""Constructs a ResNet-101 model.
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
"""
model = ResNet_imagenet(Bottleneck, [3, 4, 23, 3], **kwargs)
if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['resnet101']))
return model
def load_pretrain(model):
url = 'https://download.pytorch.org/models/resnet50-19c8e357.pth'
pretrained_dict = model_zoo.load_url(url)
model_dict = model.state_dict()
for k, v in model_dict.items():
if not "cls_fc" in k and not "domain_fc" in k:
model_dict[k] = pretrained_dict[k[k.find(".") + 1:]]
model.load_state_dict(model_dict)
return model
def resnet34(pretrained=False):
"""Constructs a ResNet-34 model.
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
"""
model = ResNet(BasicBlock, [3, 4, 6, 3])
if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['resnet34']))
return model
def resnet152(pretrained=False):
"""Constructs a ResNet-152 model.
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
"""
model = ResNet(Bottleneck, [3, 8, 36, 3])
if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['resnet152']))
return model
def _load_pretrained_model(self):
pretrain_dict = model_zoo.load_url('https://download.pytorch.org/models/resnet101-5d3b4d8f.pth')
model_dict = {}
state_dict = self.state_dict()
for k, v in pretrain_dict.items():
if k in state_dict:
model_dict[k] = v
state_dict.update(model_dict)
self.load_state_dict(state_dict)
def xresnet50_2(pretrained=False, **kwargs):
"""Constructs a XResNet-50 model.
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
"""
model = XResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
if pretrained: model.load_state_dict(model_zoo.load_url(model_urls['xresnet50']))
return model
def resnet50(pretrained=False, channel= 20, **kwargs):
model = ResNet(Bottleneck, [3, 4, 6, 3], nb_classes=101, channel=channel, **kwargs)
if pretrained:
pretrain_dict = model_zoo.load_url(model_urls['resnet50']) # modify pretrain code
model_dict = model.state_dict()
model_dict=weight_transform(model_dict, pretrain_dict, channel)
model.load_state_dict(model_dict)
return model
def main():
"""Create the model and start the evaluation process."""
args = get_arguments()
# gpu0 = args.gpu
if not os.path.exists(args.save_dir):
os.makedirs(args.save_dir)
model = Res_Deeplab(num_classes=args.num_classes)
if args.pretrained_model != None:
args.restore_from = pretrianed_models_dict[args.pretrained_model]
if args.restore_from[:4] == 'http' :
saved_state_dict = model_zoo.load_url(args.restore_from)
else:
saved_state_dict = torch.load(args.restore_from)
model.load_state_dict(saved_state_dict)
model.eval()
# model.cuda(gpu0)
testloader = data.DataLoader(VOCDataSet(args.data_dir, args.data_list, crop_size=(505, 505), mean=IMG_MEAN, scale=False, mirror=False),
batch_size=1, shuffle=False, pin_memory=True)
interp = nn.Upsample(size=(505, 505), mode='bilinear')
data_list = []
colorize = VOCColorize()
for index, batch in enumerate(testloader):
if index % 100 == 0:
print('%d processd'%(index))
image, label, size, name = batch
size = size[0].numpy()
# output = model(Variable(image, volatile=True).cuda(gpu0))
output = model(Variable(image, volatile=True).cpu())
output = interp(output).cpu().data[0].numpy()
output = output[:,:size[0],:size[1]]
gt = np.asarray(label[0].numpy()[:size[0],:size[1]], dtype=np.int)
output = output.transpose(1,2,0)
output = np.asarray(np.argmax(output, axis=2), dtype=np.int)
filename = os.path.join(args.save_dir, '{}.png'.format(name[0]))
color_file = Image.fromarray(colorize(output).transpose(1, 2, 0), 'RGB')
color_file.save(filename)
# show_all(gt, output)
data_list.append([gt.flatten(), output.flatten()])
filename = os.path.join(args.save_dir, 'result.txt')
get_iou(data_list, args.num_classes, filename)
def vgg11_bn(pretrained=False, **kwargs):
"""VGG 11-layer model (configuration "A") with batch normalization
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
"""
model = VGG(make_layers(cfg['A'], batch_norm=True), **kwargs)
if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['vgg11_bn']))
return model
def vgg13(pretrained=False, **kwargs):
"""VGG 13-layer model (configuration "B")
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
"""
model = VGG(make_layers(cfg['B']), **kwargs)
if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['vgg13']))
return model
def vgg19(input_channel=3, pretrained=False, **kwargs):
"""VGG 19-layer model (configuration "E")
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
"""
model = VGG(make_layers(cfg['E'], input_channel), **kwargs)
if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['vgg19']))
return model
def alexnet(pretrained=False, **kwargs):
r"""AlexNet model architecture from the
`"One weird trick..." <https://arxiv.org/abs/1404.5997>`_ paper.
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
"""
model = AlexNet(**kwargs)
if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['alexnet']))
return model
def cifar100(n_channel, pretrained=None):
cfg = [n_channel, n_channel, 'M', 2*n_channel, 2*n_channel, 'M', 4*n_channel, 4*n_channel, 'M', (8*n_channel, 0), 'M']
layers = make_layers(cfg, batch_norm=True)
model = CIFAR(layers, n_channel=8*n_channel, num_classes=100)
if pretrained is not None:
m = model_zoo.load_url(model_urls['cifar100'])
state_dict = m.state_dict() if isinstance(m, nn.Module) else m
assert isinstance(state_dict, (dict, OrderedDict)), type(state_dict)
model.load_state_dict(state_dict)
return model
def resnet50(pretrained=False):
"""Constructs a ResNet-50 model.
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
"""
model = ResNet(Bottleneck, [3, 4, 6, 3])
if pretrained:
model.load_state_dict(remove_fc(model_zoo.load_url(model_urls['resnet50'])))
return model
def load_pretrained(model, num_classes, settings):
assert num_classes == settings['num_classes'], \
"num_classes should be {}, but is {}".format(settings['num_classes'], num_classes)
model.load_state_dict(model_zoo.load_url(settings['url']))
model.input_space = settings['input_space']
model.input_size = settings['input_size']
model.input_range = settings['input_range']
model.mean = settings['mean']
model.std = settings['std']
return model
def resnet18(pretrained=False):
"""Constructs a ResNet-18 model.
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
"""
model = ResNet(BasicBlock, [2, 2, 2, 2])
if pretrained:
model.load_state_dict(remove_fc(model_zoo.load_url(model_urls['resnet18'])))
return model
def get_encoder(name, encoder_weights=None):
Encoder = encoders[name]['encoder']
encoder = Encoder(**encoders[name]['params'])
encoder.out_shapes = encoders[name]['out_shapes']
if encoder_weights is not None:
settings = encoders[name]['pretrained_settings'][encoder_weights]
encoder.load_state_dict(model_zoo.load_url(settings['url']))
return encoder
def inceptionresnetv2(pretrained=True):
r"""InceptionResnetV2 model architecture from the
`"InceptionV4, Inception-ResNet..." <https://arxiv.org/abs/1602.07261>`_ paper.
Args:
pretrained ('string'): If True, returns a model pre-trained on ImageNet
"""
model = InceptionResnetV2()
if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['imagenet']))
return model
def drn_a_50(BatchNorm, pretrained=True):
model = DRN_A(Bottleneck, [3, 4, 6, 3], BatchNorm=BatchNorm)
if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['resnet50']))
return model
def test_vgg19(self):
state_dict = model_zoo.load_url(model_urls['vgg19'], progress=False)
self.run_model_test(vgg19(),
train=False,
batch_size=BATCH_SIZE,
state_dict=state_dict)
import torch
from os import path
from sys import version_info
from collections import OrderedDict
from torch.utils.model_zoo import load_url
# Download the VGG-19 model and fix the layer names
print("Downloading the VGG-19 model")
sd = load_url("https://s3-us-west-2.amazonaws.com/jcjohns-models/vgg19-d01eb7cb.pth")
map = {'classifier.1.weight':u'classifier.0.weight', 'classifier.1.bias':u'classifier.0.bias', 'classifier.4.weight':u'classifier.3.weight', 'classifier.4.bias':u'classifier.3.bias'}
sd = OrderedDict([(map[k] if k in map else k,v) for k,v in sd.items()])
torch.save(sd, path.join("models", "vgg19-d01eb7cb.pth"))
# Download the VGG-16 model and fix the layer names
print("Downloading the VGG-16 model")
sd = load_url("https://s3-us-west-2.amazonaws.com/jcjohns-models/vgg16-00b39a1b.pth")
map = {'classifier.1.weight':u'classifier.0.weight', 'classifier.1.bias':u'classifier.0.bias', 'classifier.4.weight':u'classifier.3.weight', 'classifier.4.bias':u'classifier.3.bias'}
sd = OrderedDict([(map[k] if k in map else k,v) for k,v in sd.items()])
torch.save(sd, path.join("models", "vgg16-00b39a1b.pth"))
# Download the NIN model
print("Downloading the NIN model")
if version_info[0] < 3:
import urllib
urllib.URLopener().retrieve("https://raw.githubusercontent.com/ProGamerGov/pytorch-nin/master/nin_imagenet.pth", path.join("models", "nin_imagenet.pth"))
else:
import urllib.request
urllib.request.urlretrieve("https://raw.githubusercontent.com/ProGamerGov/pytorch-nin/master/nin_imagenet.pth", path.join("models", "nin_imagenet.pth"))
print("All models have been successfully downloaded")
def main():
"""Create the model and start the evaluation process."""
args = get_arguments()
#gpu0 = args.gpu
if not os.path.exists(args.save_dir):
os.makedirs(args.save_dir)
model = Res_Deeplab(num_classes=args.num_classes)
if args.pretrained_model != None:
args.restore_from = pretrianed_models_dict[args.pretrained_model]
if args.restore_from[:4] == 'http':
saved_state_dict = model_zoo.load_url(args.restore_from,
map_location='cpu')
else:
saved_state_dict = torch.load(args.restore_from, map_location='cpu')
model.load_state_dict(saved_state_dict)
model.eval()
#model.cuda(gpu0)
testloader = data.DataLoader(VOCDataSet(args.data_dir,
args.data_list,
crop_size=(505, 505),
mean=IMG_MEAN,
scale=False,
mirror=False),
batch_size=1,
shuffle=False,
pin_memory=True)
if version.parse(torch.__version__) >= version.parse('0.4.0'):
interp = nn.Upsample(size=(505, 505),
mode='bilinear',
align_corners=True)
else:
interp = nn.Upsample(size=(505, 505), mode='bilinear')
data_list = []
colorize = VOCColorize()
for index, batch in enumerate(testloader):
if index % 100 == 0:
print('%d processd' % (index))
image, label, size, name = batch
size = size[0].numpy()
output = model(Variable(image, volatile=True))
output = interp(output).cpu().data[0].numpy()
output = output[:, :size[0], :size[1]]
gt = np.asarray(label[0].numpy()[:size[0], :size[1]], dtype=np.int)
output = output.transpose(1, 2, 0)
output = np.asarray(np.argmax(output, axis=2), dtype=np.int)
filename = os.path.join(args.save_dir, '{}.png'.format(name[0]))
color_file = Image.fromarray(
colorize(output).transpose(1, 2, 0), 'RGB')
color_file.save(filename)
# show_all(gt, output)
data_list.append([gt.flatten(), output.flatten()])
filename = os.path.join(args.save_dir, 'result.txt')
get_iou(data_list, args.num_classes, filename)
def drn_d_54(pretrained=False, **kwargs):
model = DRN(Bottleneck, [1, 1, 3, 4, 6, 3, 1, 1], arch='D', **kwargs)
if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['drn-d-54']))
return model
def main():
print("SSSRNet3 training from scratch on VOC 160*160 patches.")
global opt, model, netContent
opt = parser.parse_args()
print(opt)
gpuid = 0
cuda = True
if cuda and not torch.cuda.is_available():
raise Exception("No GPU found, please run without --cuda")
opt.seed = random.randint(1, 10000)
print("Random Seed: ", opt.seed)
torch.manual_seed(opt.seed)
if cuda:
torch.cuda.manual_seed(opt.seed)
cudnn.benchmark = True
if opt.vgg_loss:
print('===> Loading VGG model')
netVGG = models.vgg19()
netVGG.load_state_dict(
model_zoo.load_url(
'https://download.pytorch.org/models/vgg19-dcbb9e9d.pth'))
class _content_model(nn.Module):
def __init__(self):
super(_content_model, self).__init__()
self.feature = nn.Sequential(
*list(netVGG.features.children())[:-1])
def forward(self, x):
out = self.feature(x)
return out
netContent = _content_model()
print("===> Building model")
model = Net()
print('Parameters: {}'.format(get_n_params(model)))
model_pretrained = torch.load(
'model/model_DIV2K_noBN_96_epoch_36.pth',
map_location=lambda storage, loc: storage)["model"]
finetune = False
if finetune == True:
index = 0
for (src, dst) in zip(model_pretrained.parameters(),
model.parameters()):
if index > 1:
list(model.parameters())[index].data = src.data
index = index + 1
criterion = nn.MSELoss(size_average=False)
print("===> Setting GPU")
if cuda:
model = model.cuda(gpuid)
model_pretrained = model_pretrained.cuda(gpuid)
criterion = criterion.cuda(gpuid)
if opt.vgg_loss:
netContent = netContent.cuda(gpuid)
# optionally resume from a checkpoint
if opt.resume:
if os.path.isfile(opt.resume):
print("=> loading checkpoint '{}'".format(opt.resume))
checkpoint = torch.load(opt.resume)
opt.start_epoch = checkpoint["epoch"] + 1
model.load_state_dict(checkpoint["model"].state_dict())
else:
print("=> no checkpoint found at '{}'".format(opt.resume))
# optionally copy weights from a checkpoint
if opt.pretrained:
if os.path.isfile(opt.pretrained):
print("=> loading model '{}'".format(opt.pretrained))
weights = torch.load(opt.pretrained)
model.load_state_dict(weights['model'].state_dict())
else:
print("=> no model found at '{}'".format(opt.pretrained))
print("===> Setting Optimizer")
optimizer = optim.Adam(model.parameters(), lr=opt.lr)
print("===> Training1")
#root_dir = '/tmp4/hang_data/DIV2K/DIV2K_train_320_HDF5'
root_dir = '/tmp4/hang_data/VOCdevkit/VOC2012/VOC_train_label160_HDF5'
files_num = len(os.listdir(root_dir))
for epoch in range(opt.start_epoch, opt.nEpochs + 1):
#save_checkpoint(model, epoch)
print("===> Loading datasets")
x = random.sample(os.listdir(root_dir), files_num)
for index in range(0, files_num):
train_path = os.path.join(root_dir, x[index])
print("===> Training datasets: '{}'".format(train_path))
train_set = DatasetFromHdf5(train_path)
training_data_loader = DataLoader(dataset=train_set,
num_workers=opt.threads,
batch_size=opt.batchSize,
shuffle=True)
avgloss = train(training_data_loader, optimizer, model,
model_pretrained, criterion, epoch, gpuid)
if epoch % 2 == 0:
save_checkpoint(model, epoch)
#
# 모델을 변환하기 전에 모델을 추론 모드로 바꾸기 위해서 ``torch_model.eval()`` 또는 ``torch_model.train(False)`` 를
# 호출하는 것이 중요합니다.
# 이는 dropout이나 batchnorm과 같은 연산들이 추론과 학습 모드에서 다르게 작동하기 때문에 필요합니다.
#
# 미리 학습된 가중치를 읽어옵니다
model_url = 'https://s3.amazonaws.com/pytorch/test_data/export/superres_epoch100-44c6958e.pth'
batch_size = 1 # 임의의 수
# 모델을 미리 학습된 가중치로 초기화합니다
map_location = lambda storage, loc: storage
if torch.cuda.is_available():
map_location = None
torch_model.load_state_dict(
model_zoo.load_url(model_url, map_location=map_location))
# 모델을 추론 모드로 전환합니다
torch_model.eval()
######################################################################
# 이제 Tracing이나 스크립팅을 통해서 PyTorch 모델을 변환할 수 있습니다.
# 이 튜토리얼에서는 tracing을 통해 변환된 모델을 사용하도록 하겠습니다.
# 모델을 변환하기 위해서는 ``torch.onnx.export()`` 함수를 호출합니다.
# 이 함수는 모델을 실행하여 어떤 연산자들이 출력값을 계산하는데 사용되었는지를 기록합니다.
# ``export`` 함수가 모델을 실행하기 때문에, 우리가 직접 텐서를 입력값으로 넘겨주어야 합니다.
# 이 텐서의 값은 알맞은 자료형과 모양이라면 랜덤하게 결정되어도 무방합니다.
# 특정 차원을 동적인 차원으로 지정하지 않는 이상, ONNX로 변환된 그래프의 경우 입력값의 사이즈는 모든 차원에 대해 고정됩니다.
# 예시에서는 모델이 항상 배치 사이즈 1을 사용하도록 변환하였지만, 첫번째 차원을 ``torch.onnx.export()`` 의
# ``dynamic_axes`` 인자에 동적인 차원으로 지정해주면 됩니다.
#
def resnet152(pretrained=False, **kwargs):
model = ResNet(Bottleneck, [3, 8, 36, 3], **kwargs)
if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['resnet152']))
return model
def main(args):
# random.seed(args.random_seed)
# print(args)
cudnn.enabled = True
# create network
"""
if args.generatormodel == 'TwinsAdvNet_D':
model = get_model(name=args.generatormodel, num_classes = args.num_classes)
#model = TwinsAdvNet_DL(num_classes = args.num_classes)
if args.restore_from[:4] == 'http' :
saved_state_dict = model_zoo.load_url(args.restore_from)
else:
saved_state_dict = torch.load(args.restore_from)
"""
model = get_model(name=args.generatormodel, num_classes=args.num_classes)
if args.restore_from[:4] == 'http':
saved_state_dict = model_zoo.load_url(args.restore_from)
else:
saved_state_dict = torch.load(args.restore_from)
# only copy the params that exist in current model
new_params = model.state_dict().copy()
for name, param in new_params.items():
#print(name)
if name in saved_state_dict and param.size(
) == saved_state_dict[name].size():
new_params[name].copy_(saved_state_dict[name])
# print('copy {}'.format(name))
model.load_state_dict(new_params)
model.train()
# load nets into gpu
if args.num_gpus > 1:
model = torch.nn.DataParallel(model, device_ids=range(args.num_gpus))
model.cuda()
cudnn.benchmark = True # acceleration
# init D
model_D = Discriminator(num_classes=args.num_classes)
model_D.train()
if args.num_gpus > 1:
model_D = torch.nn.DataParallel(model_D,
device_ids=range(args.num_gpus))
model_D.cuda()
#train_dataset = MITSceneParsingDataset(args.list_train, args, is_train=1)
train_dataset = MITSceneParsingDataset(args.list_train,
args,
max_iters=args.max_iters,
is_train=1)
#train_dataset_size = len(train_dataset)
#args.epoch_iters = int(train_dataset_size / (args.batch_size * args.num_gpus))
#print('train_dataset_size = {} | 1 Epoch = {} iters'.format(train_dataset_size, args.epoch_iters))
trainloader = data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=int(args.workers),
pin_memory=True,
drop_last=True)
val_dataset = MITSceneParsingDataset(args.list_val,
args,
max_sample=args.num_val,
is_train=0)
val_loader = data.DataLoader(
val_dataset,
batch_size=args.batch_size,
shuffle=False, # False
num_workers=8,
pin_memory=True,
drop_last=True)
trainloader_iter = enumerate(trainloader)
# loss / bilinear upsampling
#bce_loss = BCEWithLogitsLoss2d()
bce_loss = torch.nn.BCEWithLogitsLoss() # only 0, 1
#crit = nn.NLLLoss2d(ignore_index=-1) # ade20k
crit = CrossEntropyLoss2d(ignore_index=-1)
# trainloader_iter = iter(trainloader)
# implement model.optim_parameters(args) to handle different models' lr setting
# optimizer for segmentation networks
"""
optimizer = optim.SGD(model.optim_parameters(args),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
"""
#optimizer.zero_grad() #
"""
optimizer = optim.SGD(model.parameters(),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
# optimizer for discriminator network
optimizer_D = optim.Adam(model_D.parameters(),
lr=args.learning_rate_D,
betas=(0.9, 0.99))
#optimizer_D.zero_grad()
"""
optimizer, optimizer_D = create_optimizers(model, model_D, crit, args)
#optimizer.zero_gard()
#optimizer_D.zero_gard()
# interp = nn.Upsample(size=(384, 384), mode='bilinear')
interp = nn.Upsample(size=(args.imgSize, args.imgSize), mode='bilinear')
# interp_x1 = nn.Upsample(size=(384, 384), mode='bilinear') # G1
# interp_x2 = nn.Upsample(size=(384, 384), mode='bilinear') # G2
# main loop
history = {
split: {
'epoch': [],
'loss_pred_outputs': [],
'acc_pred_outputs': []
}
for split in ('train', 'val')
}
# initial eval
evaluate(model, val_loader, interp, crit, history, 0, args)
for epoch in range(args.start_epoch, args.num_epoches + 1):
train(model, model_D, trainloader_iter, interp, optimizer, optimizer_D,
crit, bce_loss, history, epoch, args)
if epoch % args.eval_epoch == 0:
evaluate(model, val_loader, interp, crit, history, epoch, args)
end_time = timeit.default_timer()
print(' running time(s): [{0:.4f} seconds]'.format(
(end_time - start_time)))
def resnet34(pretrained=True):
model = ResNet(BasicBlock, [3, 4, 6, 3])
if pretrained:
load_weights_sequential(model,
model_zoo.load_url(model_urls['resnet34']))
return model
def main():
h, w = map(int, args.input_size.split(','))
input_size = (h, w)
cudnn.enabled = True
# gpu = args.gpu
# create network
model = Res_Deeplab(num_classes=args.num_classes)
# load pretrained parameters
if args.restore_from[:4] == 'http' :
saved_state_dict = model_zoo.load_url(args.restore_from)
else:
saved_state_dict = torch.load(args.restore_from)
# only copy the params that exist in current model (caffe-like)
new_params = model.state_dict().copy()
for name, param in new_params.items():
print(name)
if name in saved_state_dict and param.size() == saved_state_dict[name].size():
new_params[name].copy_(saved_state_dict[name])
print('copy {}'.format(name))
model.load_state_dict(new_params)
model.train()
model.cuda(args.gpu)
cudnn.benchmark = True
# init D
model_D = FCDiscriminator(num_classes=args.num_classes)
if args.restore_from_D is not None:
model_D.load_state_dict(torch.load(args.restore_from_D))
model_D.train()
model_D.cuda(args.gpu)
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
train_dataset = VOCDataSet(args.data_dir, args.data_list, crop_size=input_size,
scale=args.random_scale, mirror=args.random_mirror, mean=IMG_MEAN)
train_dataset_size = len(train_dataset)
train_gt_dataset = VOCGTDataSet(args.data_dir, args.data_list, crop_size=input_size,
scale=args.random_scale, mirror=args.random_mirror, mean=IMG_MEAN)
if args.partial_data is None:
trainloader = data.DataLoader(train_dataset,
batch_size=args.batch_size, shuffle=True, num_workers=5, pin_memory=True)
trainloader_gt = data.DataLoader(train_gt_dataset,
batch_size=args.batch_size, shuffle=True, num_workers=5, pin_memory=True)
else:
#sample partial data
partial_size = int(args.partial_data * train_dataset_size)
if args.partial_id is not None:
train_ids = pickle.load(open(args.partial_id))
print('loading train ids from {}'.format(args.partial_id))
else:
train_ids = list(range(train_dataset_size))
np.random.shuffle(train_ids)
pickle.dump(train_ids, open(osp.join(args.snapshot_dir, 'train_id.pkl'), 'wb'))
train_sampler = data.sampler.SubsetRandomSampler(train_ids[:partial_size])
train_remain_sampler = data.sampler.SubsetRandomSampler(train_ids[partial_size:])
train_gt_sampler = data.sampler.SubsetRandomSampler(train_ids[:partial_size])
trainloader = data.DataLoader(train_dataset,
batch_size=args.batch_size, sampler=train_sampler, num_workers=3, pin_memory=True)
trainloader_remain = data.DataLoader(train_dataset,
batch_size=args.batch_size, sampler=train_remain_sampler, num_workers=3, pin_memory=True)
trainloader_gt = data.DataLoader(train_gt_dataset,
batch_size=args.batch_size, sampler=train_gt_sampler, num_workers=3, pin_memory=True)
# trainloader_remain_iter = enumerate(trainloader_remain)
trainloader_remain_iter = iter(trainloader_remain)
# trainloader_iter = enumerate(trainloader)
# trainloader_gt_iter = enumerate(trainloader_gt)
trainloader_iter = iter(trainloader)
trainloader_gt_iter = iter(trainloader_gt)
# implement model.optim_parameters(args) to handle different models' lr setting
# optimizer for segmentation network
optimizer = optim.SGD(model.optim_parameters(args),
lr=args.learning_rate, momentum=args.momentum,weight_decay=args.weight_decay)
optimizer.zero_grad()
# optimizer for discriminator network
optimizer_D = optim.Adam(model_D.parameters(), lr=args.learning_rate_D, betas=(0.9,0.99))
optimizer_D.zero_grad()
# loss/ bilinear upsampling
bce_loss = BCEWithLogitsLoss2d()
interp = nn.Upsample(size=(input_size[1], input_size[0]), mode='bilinear')
if version.parse(torch.__version__) >= version.parse('0.4.0'):
interp = nn.Upsample(size=(input_size[1], input_size[0]), mode='bilinear', align_corners=True)
else:
interp = nn.Upsample(size=(input_size[1], input_size[0]), mode='bilinear')
# labels for adversarial training
pred_label = 0
gt_label = 1
for i_iter in range(args.num_steps):
loss_seg_value = 0
loss_adv_pred_value = 0
loss_D_value = 0
loss_semi_value = 0
loss_semi_adv_value = 0
optimizer.zero_grad()
adjust_learning_rate(optimizer, i_iter)
optimizer_D.zero_grad()
adjust_learning_rate_D(optimizer_D, i_iter)
for sub_i in range(args.iter_size):
# train G
# don't accumulate grads in D
for param in model_D.parameters():
param.requires_grad = False
# do semi first
if (args.lambda_semi > 0 or args.lambda_semi_adv > 0 ) and i_iter >= args.semi_start_adv :
try:
# _, batch = trainloader_remain_iter.next()
batch = next(trainloader_remain_iter)
except:
trainloader_remain_iter = iter(trainloader_remain)
batch = next(trainloader_remain_iter)
# only access to img
images, _, _, _ = batch
images = Variable(images).cuda(args.gpu)
pred = interp(model(images))
pred_remain = pred.detach()
D_out = interp(model_D(F.softmax(pred)))
D_out_sigmoid = F.sigmoid(D_out).data.cpu().numpy().squeeze(axis=1)
ignore_mask_remain = np.zeros(D_out_sigmoid.shape).astype(np.bool)
loss_semi_adv = args.lambda_semi_adv * bce_loss(D_out, make_D_label(gt_label, ignore_mask_remain))
loss_semi_adv = loss_semi_adv/args.iter_size
##loss_semi_adv.backward()
# loss_semi_adv_value += loss_semi_adv.data.cpu().numpy()[0]/args.lambda_semi_adv
loss_semi_adv_value += loss_semi_adv.item() / args.lambda_semi_adv
if args.lambda_semi <= 0 or i_iter < args.semi_start:
loss_semi_adv.backward()
loss_semi_value = 0
else:
# produce ignore mask
semi_ignore_mask = (D_out_sigmoid < args.mask_T)
semi_gt = pred.data.cpu().numpy().argmax(axis=1)
semi_gt[semi_ignore_mask] = 255
semi_ratio = 1.0 - float(semi_ignore_mask.sum())/semi_ignore_mask.size
print('semi ratio: {:.4f}'.format(semi_ratio))
if semi_ratio == 0.0:
loss_semi_value += 0
else:
semi_gt = torch.FloatTensor(semi_gt)
loss_semi = args.lambda_semi * loss_calc(pred, semi_gt, args.gpu)
loss_semi = loss_semi/args.iter_size
# loss_semi_value += loss_semi.data.cpu().numpy()[0]/args.lambda_semi
loss_semi_value += loss_semi.item() / args.lambda_semi
loss_semi += loss_semi_adv
loss_semi.backward()
else:
loss_semi = None
loss_semi_adv = None
# train with source
try:
# _, batch = trainloader_iter.next()
batch = next(trainloader_iter)
except:
trainloader_iter = iter(trainloader)
batch = next(trainloader_iter)
images, labels, _, _ = batch
images = Variable(images).cuda(args.gpu)
ignore_mask = (labels.numpy() == 255)
pred = interp(model(images)) # [5,21,41,41] -> [5, 21, 321, 321]
loss_seg = loss_calc(pred, labels, args.gpu)
D_out = interp(model_D(F.softmax(pred))) # [5, 1, 10, 10] -> [5, 1, 321, 321]
loss_adv_pred = bce_loss(D_out, make_D_label(gt_label, ignore_mask))
loss = loss_seg + args.lambda_adv_pred * loss_adv_pred
# proper normalization
loss = loss/args.iter_size
loss.backward()
# loss_seg_value += loss_seg.data.cpu().numpy()[0]/args.iter_size
# loss_adv_pred_value += loss_adv_pred.data.cpu().numpy()[0]/args.iter_size
loss_seg_value += loss_seg.item() / args.iter_size
loss_adv_pred_value += loss_adv_pred.item() / args.iter_size
# train D
# bring back requires_grad
for param in model_D.parameters():
param.requires_grad = True
# train with pred
pred = pred.detach()
if args.D_remain:
try:
pred_remain
pred = torch.cat((pred, pred_remain), 0)
ignore_mask = np.concatenate((ignore_mask, ignore_mask_remain), axis=0)
except NameError:
print("No pred_remain during training D")
# pred = torch.cat((pred, pred_remain), 0)
# ignore_mask = np.concatenate((ignore_mask,ignore_mask_remain), axis = 0)
D_out = interp(model_D(F.softmax(pred)))
loss_D = bce_loss(D_out, make_D_label(pred_label, ignore_mask))
loss_D = loss_D/args.iter_size/2
loss_D.backward()
# loss_D_value += loss_D.data.cpu().numpy()[0]
loss_D_value += loss_D.item()
# train with gt
# get gt labels
try:
# _, batch = trainloader_gt_iter.next()
batch = next(trainloader_gt_iter)
except:
trainloader_gt_iter = iter(trainloader_gt)
batch = next(trainloader_gt_iter)
_, labels_gt, _, _ = batch
D_gt_v = Variable(one_hot(labels_gt)).cuda(args.gpu)
ignore_mask_gt = (labels_gt.numpy() == 255)
D_out = interp(model_D(D_gt_v))
loss_D = bce_loss(D_out, make_D_label(gt_label, ignore_mask_gt))
loss_D = loss_D/args.iter_size/2
loss_D.backward()
# loss_D_value += loss_D.data.cpu().numpy()[0]
loss_D_value += loss_D.item()
optimizer.step()
optimizer_D.step()
print('exp = {}'.format(args.snapshot_dir))
print('iter = {0:8d}/{1:8d}, loss_seg = {2:.3f}, loss_adv_p = {3:.3f}, loss_D = {4:.3f}, loss_semi = {5:.3f}, loss_semi_adv = {6:.3f}'.format(i_iter, args.num_steps, loss_seg_value, loss_adv_pred_value, loss_D_value, loss_semi_value, loss_semi_adv_value))
if i_iter >= args.num_steps-1:
print('save model ...')
torch.save(model.state_dict(),osp.join(args.snapshot_dir, 'VOC_'+str(args.num_steps)+'.pth'))
torch.save(model_D.state_dict(),osp.join(args.snapshot_dir, 'VOC_'+str(args.num_steps)+'_D.pth'))
break
if i_iter % args.save_pred_every == 0 and i_iter!=0:
print('taking snapshot ...')
torch.save(model.state_dict(),osp.join(args.snapshot_dir, 'VOC_'+str(i_iter)+'.pth'))
torch.save(model_D.state_dict(),osp.join(args.snapshot_dir, 'VOC_'+str(i_iter)+'_D.pth'))
end = timeit.default_timer()
print(end-start,'seconds')
def main():
logger.auto_set_dir()
global args
parser = argparse.ArgumentParser()
parser = argparse.ArgumentParser()
parser.add_argument('--dataroot',
default='/home/hutao/lab/pytorchgo/example/ROAD/data',
help='Path to source dataset')
parser.add_argument('--batchSize',
type=int,
default=1,
help='input batch size')
parser.add_argument('--max_epoch',
type=int,
default=max_epoch,
help='Number of training iterations')
parser.add_argument('--optimizer',
type=str,
default='Adam',
help='Optimizer to use | SGD, Adam')
parser.add_argument('--lr',
type=float,
default=base_lr,
help='learning rate')
parser.add_argument('--momentum',
type=float,
default=0.99,
help='Momentum for SGD')
parser.add_argument('--beta1',
type=float,
default=0.9,
help='beta1 for adam. default=0.5')
parser.add_argument('--weight_decay',
type=float,
default=0.0005,
help='Weight decay')
parser.add_argument('--model', type=str, default='vgg16')
parser.add_argument('--gpu', type=int, default=2)
args = parser.parse_args()
print(args)
gpu = args.gpu
os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu)
cuda = torch.cuda.is_available()
torch.manual_seed(1337)
if cuda:
logger.info("random seed 1337")
torch.cuda.manual_seed(1337)
# Defining data loaders
kwargs = {
'num_workers': 4,
'pin_memory': True,
'drop_last': True
} if cuda else {}
train_loader = torch.utils.data.DataLoader(torchfcn.datasets.SYNTHIA(
'SYNTHIA',
args.dataroot,
split='train',
transform=True,
image_size=image_size),
batch_size=args.batchSize,
shuffle=True,
**kwargs)
val_loader = torch.utils.data.DataLoader(torchfcn.datasets.CityScapes(
'cityscapes',
args.dataroot,
split='val',
transform=True,
image_size=image_size),
batch_size=1,
shuffle=False)
target_loader = torch.utils.data.DataLoader(torchfcn.datasets.CityScapes(
'cityscapes',
args.dataroot,
split='train',
transform=True,
image_size=image_size),
batch_size=args.batchSize,
shuffle=True)
if cuda:
torch.set_default_tensor_type('torch.cuda.FloatTensor')
if args.model == "vgg16":
model = origin_model = torchfcn.models.Seg_model(n_class=class_num)
vgg16 = torchfcn.models.VGG16(pretrained=True)
model.copy_params_from_vgg16(vgg16)
model_fix = torchfcn.models.Seg_model(n_class=class_num)
model_fix.copy_params_from_vgg16(vgg16)
elif args.model == "deeplabv2": # TODO may have problem!
model = origin_model = torchfcn.models.Res_Deeplab(
num_classes=class_num, image_size=image_size)
saved_state_dict = model_zoo.load_url(Deeplabv2_restore_from)
new_params = model.state_dict().copy()
for i in saved_state_dict:
# Scale.layer5.conv2d_list.3.weight
i_parts = i.split('.')
# print i_parts
if not class_num == 19 or not i_parts[1] == 'layer5':
new_params['.'.join(i_parts[1:])] = saved_state_dict[i]
# print i_parts
model.load_state_dict(new_params)
model_fix = torchfcn.models.Res_Deeplab(num_classes=class_num,
image_size=image_size)
model_fix.load_state_dict(new_params)
else:
raise ValueError("only support vgg16, deeplabv2!")
for param in model_fix.parameters():
param.requires_grad = False
netD = torchfcn.models.Domain_classifer(n_class=class_num)
netD.apply(weights_init)
model_summary([model, netD])
if cuda:
model = model.cuda()
netD = netD.cuda()
# Defining optimizer
if args.optimizer == 'SGD':
raise ValueError("SGD is not prepared well..")
optim = torch.optim.SGD([
{
'params': get_parameters(model, bias=False)
},
{
'params': get_parameters(model, bias=True),
'lr': args.lr * 2,
'weight_decay': args.weight_decay
},
],
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
elif args.optimizer == 'Adam':
if args.model == "vgg16":
optim = torch.optim.Adam([
{
'params': get_parameters(model, bias=False),
'weight_decay': args.weight_decay
},
{
'params': get_parameters(model, bias=True),
'lr': args.lr * 2,
'weight_decay': args.weight_decay
},
],
lr=args.lr,
betas=(args.beta1, 0.999))
elif args.model == "deeplabv2":
optim = torch.optim.Adam(origin_model.optim_parameters(args.lr),
lr=args.lr,
betas=(args.beta1, 0.999),
weight_decay=args.weight_decay)
else:
raise
else:
raise ValueError('Invalid optmizer argument. Has to be SGD or Adam')
optimD = torch.optim.Adam(netD.parameters(),
lr=dis_lr,
weight_decay=args.weight_decay,
betas=(0.7, 0.999))
optimizer_summary([optim, optimD])
trainer = MyTrainer_ROAD(cuda=cuda,
model=model,
model_fix=model_fix,
netD=netD,
optimizer=optim,
optimizerD=optimD,
train_loader=train_loader,
target_loader=target_loader,
val_loader=val_loader,
batch_size=args.batchSize,
image_size=image_size,
loss_print_interval=LOSS_PRINT_INTERVAL)
trainer.epoch = 0
trainer.iteration = 0
trainer.train()
def load_pretrained(model,
cfg=None,
num_classes=1000,
in_chans=3,
filter_fn=None,
strict=True):
if cfg is None:
cfg = getattr(model, 'default_cfg')
if cfg is None or 'url' not in cfg or not cfg['url']:
_logger.warning(
"Pretrained model URL is invalid, using random initialization.")
return
state_dict = model_zoo.load_url(cfg['url'],
progress=False,
map_location='cpu')
if filter_fn is not None:
state_dict = filter_fn(state_dict)
if in_chans == 1:
conv1_name = cfg['first_conv']
_logger.info(
'Converting first conv (%s) pretrained weights from 3 to 1 channel'
% conv1_name)
conv1_weight = state_dict[conv1_name + '.weight']
# Some weights are in torch.half, ensure it's float for sum on CPU
conv1_type = conv1_weight.dtype
conv1_weight = conv1_weight.float()
O, I, J, K = conv1_weight.shape
if I > 3:
assert conv1_weight.shape[1] % 3 == 0
# For models with space2depth stems
conv1_weight = conv1_weight.reshape(O, I // 3, 3, J, K)
conv1_weight = conv1_weight.sum(dim=2, keepdim=False)
else:
conv1_weight = conv1_weight.sum(dim=1, keepdim=True)
conv1_weight = conv1_weight.to(conv1_type)
state_dict[conv1_name + '.weight'] = conv1_weight
elif in_chans != 3:
conv1_name = cfg['first_conv']
conv1_weight = state_dict[conv1_name + '.weight']
conv1_type = conv1_weight.dtype
conv1_weight = conv1_weight.float()
O, I, J, K = conv1_weight.shape
if I != 3:
_logger.warning(
'Deleting first conv (%s) from pretrained weights.' %
conv1_name)
del state_dict[conv1_name + '.weight']
strict = False
else:
# NOTE this strategy should be better than random init, but there could be other combinations of
# the original RGB input layer weights that'd work better for specific cases.
_logger.info('Repeating first conv (%s) weights in channel dim.' %
conv1_name)
repeat = int(math.ceil(in_chans / 3))
conv1_weight = conv1_weight.repeat(1, repeat, 1,
1)[:, :in_chans, :, :]
conv1_weight *= (3 / float(in_chans))
conv1_weight = conv1_weight.to(conv1_type)
state_dict[conv1_name + '.weight'] = conv1_weight
classifier_name = cfg['classifier']
if num_classes == 1000 and cfg['num_classes'] == 1001:
# special case for imagenet trained models with extra background class in pretrained weights
classifier_weight = state_dict[classifier_name + '.weight']
state_dict[classifier_name + '.weight'] = classifier_weight[1:]
classifier_bias = state_dict[classifier_name + '.bias']
state_dict[classifier_name + '.bias'] = classifier_bias[1:]
elif num_classes != cfg['num_classes']:
# completely discard fully connected for all other differences between pretrained and created model
del state_dict[classifier_name + '.weight']
del state_dict[classifier_name + '.bias']
strict = False
model.load_state_dict(state_dict, strict=strict)
def main():
tag = 0
h, w = map(int, args.input_size.split(','))
input_size = (h, w)
cudnn.enabled = True
# create network
model = Res_Deeplab(num_classes=args.num_classes)
model = nn.DataParallel(model)
model.cuda()
# load pretrained parameters
if args.restore_from[:4] == 'http':
saved_state_dict = model_zoo.load_url(args.restore_from)
else:
saved_state_dict = torch.load(args.restore_from)
# only copy the params that exist in current model (caffe-like)
saved_state_dict = torch.load(
'/data1/wyc/AdvSemiSeg/snapshots/VOC_t_baseline_1adv_mul_20000.pth')
new_params = model.state_dict().copy()
for name, param in new_params.items():
print(name)
if name in saved_state_dict and param.size(
) == saved_state_dict[name].size():
new_params[name].copy_(saved_state_dict[name])
print('copy {}'.format(name))
else:
print 123456
model.load_state_dict(new_params)
model.train()
cudnn.benchmark = True
# init D
model_D = Discriminator2(num_classes=args.num_classes)
if args.restore_from_D is not None:
model_D.load_state_dict(torch.load(args.restore_from_D))
model_D = nn.DataParallel(model_D)
model_D.train()
model_D.cuda()
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
train_dataset = VOCDataSet(args.data_dir,
args.data_list,
crop_size=input_size,
scale=args.random_scale,
mirror=args.random_mirror,
mean=IMG_MEAN)
train_dataset_size = len(train_dataset)
train_gt_dataset = VOCGTDataSet(args.data_dir,
args.data_list,
crop_size=input_size,
scale=args.random_scale,
mirror=args.random_mirror,
mean=IMG_MEAN)
if args.partial_data == 0:
trainloader = data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=5,
pin_memory=True)
trainloader_gt = data.DataLoader(train_gt_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=5,
pin_memory=True)
else:
#sample partial data
partial_size = int(args.partial_data * train_dataset_size)
trainloader_iter = enumerate(trainloader)
trainloader_gt_iter = enumerate(trainloader_gt)
# implement model.optim_parameters(args) to handle different models' lr setting
# optimizer for segmentation network
optimizer = optim.SGD(model.module.optim_parameters(args),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
optimizer.zero_grad()
# optimizer for discriminator network
optimizer_D = optim.Adam(model_D.parameters(),
lr=args.learning_rate_D,
betas=(0.9, 0.99))
optimizer_D.zero_grad()
# loss/ bilinear upsampling
bce_loss = torch.nn.BCELoss()
interp = nn.Upsample(size=(input_size[1], input_size[0]), mode='bilinear')
if version.parse(torch.__version__) >= version.parse('0.4.0'):
interp = nn.Upsample(size=(input_size[1], input_size[0]),
mode='bilinear',
align_corners=True)
else:
interp = nn.Upsample(size=(input_size[1], input_size[0]),
mode='bilinear')
# labels for adversarial training
pred_label = 0
gt_label = 1
for i_iter in range(args.num_steps):
loss_seg_value = 0
loss_adv_pred_value = 0
loss_D_value = 0
loss_semi_value = 0
loss_semi_adv_value = 0
optimizer.zero_grad()
adjust_learning_rate(optimizer, i_iter)
optimizer_D.zero_grad()
adjust_learning_rate_D(optimizer_D, i_iter)
for sub_i in range(args.iter_size):
# train G
# don't accumulate grads in D
for param in model_D.parameters():
param.requires_grad = False
# train with source
try:
_, batch = trainloader_iter.next()
except:
trainloader_iter = enumerate(trainloader)
_, batch = trainloader_iter.next()
images, labels, _, _ = batch
images = Variable(images).cuda()
ignore_mask = (labels.numpy() == 255)
pred = interp(model(images))
loss_seg = loss_calc(pred, labels)
pred01 = F.softmax(pred, dim=1)
pred_re = F.softmax(pred, dim=1).repeat(1, 3, 1, 1)
indices_1 = torch.index_select(
images, 1,
Variable(torch.LongTensor([0])).cuda())
indices_2 = torch.index_select(
images, 1,
Variable(torch.LongTensor([1])).cuda())
indices_3 = torch.index_select(
images, 1,
Variable(torch.LongTensor([2])).cuda())
img_re = torch.cat([
indices_1.repeat(1, 21, 1, 1),
indices_2.repeat(1, 21, 1, 1),
indices_3.repeat(1, 21, 1, 1),
], 1)
mul_img = pred_re * img_re
for i_l in range(labels.shape[0]):
label_set = np.unique(labels[i_l]).tolist()
for ls in label_set:
if ls != 0 and ls != 255:
ls = int(ls)
img_p = torch.cat([
mul_img[i_l][ls].unsqueeze(0).unsqueeze(0),
mul_img[i_l][ls + 21].unsqueeze(0).unsqueeze(0),
mul_img[i_l][ls + 21 +
21].unsqueeze(0).unsqueeze(0)
], 1)
imgs = img_p.squeeze()
imgs = imgs.transpose(0, 1)
imgs = imgs.transpose(1, 2)
imgs = imgs.data.cpu().numpy()
img_ori = images[0]
img_ori = img_ori.squeeze()
img_ori = img_ori.transpose(0, 1)
img_ori = img_ori.transpose(1, 2)
img_ori = img_ori.data.cpu().numpy()
pred_ori = pred01[0][ls]
pred_ori = pred_ori.data.cpu().numpy()
pred_ori = pred_ori[:, :, np.newaxis] * 255.0
# print pred_ori.shape
if tag == 0:
print ls
cv2.imwrite('/data1/wyc/1.png', imgs)
cv2.imwrite('/data1/wyc/2.png', img_ori)
cv2.imwrite('/data1/wyc/3.png', pred_ori)
tag = 1
D_out = model_D(mul_img)
loss_adv_pred = bce_loss(D_out, make_D_label(gt_label, D_out))
loss = loss_seg + args.lambda_adv_pred * loss_adv_pred
# proper normalization
loss = loss / args.iter_size
loss.backward()
loss_seg_value += loss_seg.data.cpu().numpy()[0] / args.iter_size
loss_adv_pred_value += loss_adv_pred.data.cpu().numpy(
)[0] / args.iter_size
# train D
# bring back requires_grad
for param in model_D.parameters():
param.requires_grad = True
# train with pred
pred = pred.detach()
pred_re2 = F.softmax(pred, dim=1).repeat(1, 3, 1, 1)
mul_img2 = pred_re2 * img_re
D_out = model_D(mul_img2)
loss_D = bce_loss(D_out, make_D_label(pred_label, D_out))
loss_D = loss_D / args.iter_size / 2
loss_D.backward()
loss_D_value += loss_D.data.cpu().numpy()[0]
# train with gt
# get gt labels
try:
_, batch = trainloader_gt_iter.next()
except:
trainloader_gt_iter = enumerate(trainloader_gt)
_, batch = trainloader_gt_iter.next()
img_gt, labels_gt, _, _ = batch
img_gt = Variable(img_gt).cuda()
D_gt_v = Variable(one_hot(labels_gt)).cuda()
ignore_mask_gt = (labels_gt.numpy() == 255)
pred_re3 = D_gt_v.repeat(1, 3, 1, 1)
indices_1 = torch.index_select(
img_gt, 1,
Variable(torch.LongTensor([0])).cuda())
indices_2 = torch.index_select(
img_gt, 1,
Variable(torch.LongTensor([1])).cuda())
indices_3 = torch.index_select(
img_gt, 1,
Variable(torch.LongTensor([2])).cuda())
img_re3 = torch.cat([
indices_1.repeat(1, 21, 1, 1),
indices_2.repeat(1, 21, 1, 1),
indices_3.repeat(1, 21, 1, 1),
], 1)
mul_img3 = pred_re3 * img_re3
D_out = model_D(mul_img3)
loss_D = bce_loss(D_out, make_D_label(gt_label, D_out))
loss_D = loss_D / args.iter_size / 2
loss_D.backward()
loss_D_value += loss_D.data.cpu().numpy()[0]
optimizer.step()
optimizer_D.step()
print('exp = {}'.format(args.snapshot_dir))
print(
'iter = {0:8d}/{1:8d}, loss_seg = {2:.3f}, loss_adv_p = {3:.3f}, loss_D = {4:.3f}, loss_semi = {5:.3f}, loss_semi_adv = {6:.3f}'
.format(i_iter, args.num_steps, loss_seg_value,
loss_adv_pred_value, loss_D_value, loss_semi_value,
loss_semi_adv_value))
if i_iter >= args.num_steps - 1:
print('save model ...')
torch.save(
model.state_dict(),
osp.join(
args.snapshot_dir, 'VOC_' +
os.path.abspath(__file__).split('/')[-1].split('.')[0] +
'_' + str(args.num_steps) + '.pth'))
#torch.save(model_D.state_dict(),osp.join(args.snapshot_dir, 'VOC_'+os.path.abspath(__file__).split('/')[-1].split('.')[0]+'_'+str(args.num_steps)+'_D.pth'))
break
if i_iter % 100 == 0 and i_iter != 0:
print('taking snapshot ...')
torch.save(
model.state_dict(),
osp.join(
args.snapshot_dir, 'VOC_' +
os.path.abspath(__file__).split('/')[-1].split('.')[0] +
'_' + str(i_iter) + '.pth'))
#torch.save(model_D.state_dict(),osp.join(args.snapshot_dir, 'VOC_'+os.path.abspath(__file__).split('/')[-1].split('.')[0]+'_'+str(i_iter)+'_D.pth'))
end = timeit.default_timer()
print(end - start, 'seconds')
###################################################################### # Ordinarily, you would now train this model; however, for this tutorial, # we will instead download some pre-trained weights. Note that this model # was not trained fully for good accuracy and is used here for # demonstration purposes only. # # Load pretrained model weights model_url = 'https://s3.amazonaws.com/pytorch/test_data/export/superres_epoch100-44c6958e.pth' batch_size = 1 # just a random number # Initialize model with the pretrained weights map_location = lambda storage, loc: storage if torch.cuda.is_available(): map_location = None torch_model.load_state_dict(model_zoo.load_url(model_url, map_location=map_location)) # set the train mode to false since we will only run the forward pass. torch_model.train(False) ###################################################################### # Exporting a model in PyTorch works via tracing. To export a model, you # call the ``torch.onnx._export()`` function. This will execute the model, # recording a trace of what operators are used to compute the outputs. # Because ``_export`` runs the model, we need provide an input tensor # ``x``. The values in this tensor are not important; it can be an image # or a random tensor as long as it is the right size. # # To learn more details about PyTorch's export interface, check out the # `torch.onnx documentation <http://pytorch.org/docs/master/onnx.html>`__.
data_layer = RoIDataLayer(roidb, imdb.num_classes) # generate random permutation, '_get_next_minibatch', '_get_next_minibatch_inds'
# Create network and initialize
net = WSDDN(classes=imdb.classes, debug=_DEBUG)
net.features = torch.nn.DataParallel(net.features)
net.roi_pool = torch.nn.DataParallel(net.roi_pool)
net.classifier = torch.nn.DataParallel(net.classifier)
net.score_cls = torch.nn.DataParallel(net.score_cls)
net.score_det = torch.nn.DataParallel(net.score_det)
network.weights_normal_init(net, dev=0.001)
if os.path.exists('pretrained_alexnet.pkl'):
pret_net = pkl.load(open('pretrained_alexnet.pkl','r'))
else:
pret_net = model_zoo.load_url('https://download.pytorch.org/models/alexnet-owt-4df8aa71.pth')
pkl.dump(pret_net, open('pretrained_alexnet.pkl','wb'), pkl.HIGHEST_PROTOCOL)
own_state = net.state_dict() # net.state_dict(),keys =
for name, param in pret_net.items():
if name not in own_state:
continue
if isinstance(param, Parameter):
param = param.data
try:
own_state[name].copy_(param)
print('Copied {}'.format(name))
except:
print('Did not find {}'.format(name))
continue
if 'features' in name:
name = name.replace('features.','features.module.')
def load_pretrained(model, url):
state_dict = model_zoo.load_url(url)
state_dict = {k[7:]: v for k, v in state_dict.items() if 'module.' in k}
model.load_state_dict(state_dict, strict=False)
def xresnet(expansion, n_layers, name, pretrained=False, **kwargs):
model = XResNet(expansion, n_layers, **kwargs)
if pretrained: model.load_state_dict(model_zoo.load_url(model_urls[name]))
return model
def drn_d_107(pretrained=False, **kwargs):
model = DRN(Bottleneck, [1, 1, 3, 4, 23, 3, 2, 2], arch='D', **kwargs)
if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['drn-d-107']))
return model
def main():
"""Create the model and start the training."""
if RESTART:
args.snapshot_dir = RESTART_FROM
else:
args.snapshot_dir = generate_snapshot_name(args)
args_dict = vars(args)
import json
###### load args for restart ######
if RESTART:
# pdb.set_trace()
args_dict_file = args.snapshot_dir + '/args_dict_{}.json'.format(
RESTART_ITER)
with open(args_dict_file) as f:
args_dict_last = json.load(f)
for arg in args_dict:
args_dict[arg] = args_dict_last[arg]
###### load args for restart ######
device = torch.device("cuda" if not args.cpu else "cpu")
w, h = map(int, args.input_size.split(','))
input_size = (w, h)
w, h = map(int, args.input_size_target.split(','))
input_size_target = (w, h)
cudnn.enabled = True
cudnn.benchmark = True
if args.model == 'DeepLab':
model = DeeplabMulti(num_classes=args.num_classes)
model_D = FCDiscriminatorTest(num_classes=2 * args.num_classes).to(device)
#### restore model_D and model
if RESTART:
# pdb.set_trace()
# model parameters
restart_from_model = args.restart_from + 'GTA5_{}.pth'.format(
RESTART_ITER)
saved_state_dict = torch.load(restart_from_model)
model.load_state_dict(saved_state_dict)
# model_D parameters
restart_from_D = args.restart_from + 'GTA5_{}_D.pth'.format(
RESTART_ITER)
saved_state_dict = torch.load(restart_from_D)
model_D.load_state_dict(saved_state_dict)
#### model_D1, D2 are randomly initialized, model is pre-trained ResNet on ImageNet
else:
# model parameters
if args.restore_from[:4] == 'http':
saved_state_dict = model_zoo.load_url(args.restore_from)
else:
saved_state_dict = torch.load(args.restore_from)
new_params = model.state_dict().copy()
for i in saved_state_dict:
# Scale.layer5.conv2d_list.3.weight
i_parts = i.split('.')
# print i_parts
if not args.num_classes == 19 or not i_parts[1] == 'layer5':
new_params['.'.join(i_parts[1:])] = saved_state_dict[i]
# print i_parts
model.load_state_dict(new_params)
model.train()
model.to(device)
model_D.train()
model_D.to(device)
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
trainloader = data.DataLoader(GTA5DataSet(args.data_dir,
args.data_list,
max_iters=args.num_steps *
args.iter_size * args.batch_size,
crop_size=input_size,
scale=args.random_scale,
mirror=args.random_mirror,
mean=IMG_MEAN),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
trainloader_iter = enumerate(trainloader)
targetloader = data.DataLoader(cityscapesDataSet(
args.data_dir_target,
args.data_list_target,
max_iters=args.num_steps * args.iter_size * args.batch_size,
crop_size=input_size_target,
scale=False,
mirror=args.random_mirror,
mean=IMG_MEAN,
set=args.set),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
targetloader_iter = enumerate(targetloader)
# implement model.optim_parameters(args) to handle different models' lr setting
optimizer = optim.SGD(model.optim_parameters(args),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
optimizer.zero_grad()
optimizer_D = optim.Adam(model_D.parameters(),
lr=args.learning_rate_D,
betas=(0.9, 0.99))
optimizer_D.zero_grad()
"""
optimizer_D1 = optim.Adam(model_D1.parameters(), lr=args.learning_rate_D, betas=(0.9, 0.99))
optimizer_D1.zero_grad()
optimizer_D2 = optim.Adam(model_D2.parameters(), lr=args.learning_rate_D, betas=(0.9, 0.99))
optimizer_D2.zero_grad()
"""
if args.gan == 'Vanilla':
bce_loss = torch.nn.BCEWithLogitsLoss()
elif args.gan == 'LS':
bce_loss = torch.nn.MSELoss()
seg_loss = torch.nn.CrossEntropyLoss(ignore_index=255)
interp = nn.Upsample(size=(input_size[1], input_size[0]),
mode='bilinear',
align_corners=True)
interp_target = nn.Upsample(size=(input_size_target[1],
input_size_target[0]),
mode='bilinear',
align_corners=True)
# labels for adversarial training
source_label = 0
target_label = 1
# set up tensor board
if not os.path.exists(args.log_dir):
os.makedirs(args.log_dir)
writer = SummaryWriter(args.log_dir)
for i_iter in range(args.num_steps):
# pdb.set_trace()
loss_seg_value1 = 0
loss_seg_value2 = 0
adv_loss_value = 0
d_loss_value = 0
optimizer.zero_grad()
adjust_learning_rate(optimizer, i_iter)
optimizer_D.zero_grad()
adjust_learning_rate(optimizer_D, i_iter)
"""
optimizer_D1.zero_grad()
optimizer_D2.zero_grad()
adjust_learning_rate_D(optimizer_D1, i_iter)
adjust_learning_rate_D(optimizer_D2, i_iter)
"""
for sub_i in range(args.iter_size):
# train G
# don't accumulate grads in D
for param in model_D.parameters():
param.requires_grad = False
"""
for param in model_D1.parameters():
param.requires_grad = False
for param in model_D2.parameters():
param.requires_grad = False
"""
# train with source
_, batch = trainloader_iter.__next__()
images, labels, _, _ = batch
images = images.to(device)
labels = labels.long().to(device)
# pdb.set_trace()
# images.size() == [1, 3, 720, 1280]
pred1, pred2 = model(images)
# pred1, pred2 size == [1, 19, 91, 161]
pred1 = interp(pred1)
pred2 = interp(pred2)
# size (1, 19, 720, 1280)
# pdb.set_trace()
# feature = nn.Softmax(dim=1)(pred1)
# softmax_out = nn.Softmax(dim=1)(pred2)
loss_seg1 = seg_loss(pred1, labels)
loss_seg2 = seg_loss(pred2, labels)
loss = loss_seg2 + args.lambda_seg * loss_seg1
# pdb.set_trace()
# proper normalization
loss = loss / args.iter_size
# TODO: uncomment
loss.backward()
loss_seg_value1 += loss_seg1.item() / args.iter_size
loss_seg_value2 += loss_seg2.item() / args.iter_size
# pdb.set_trace()
# train with target
_, batch = targetloader_iter.__next__()
for params in model_D.parameters():
params.requires_grad_(requires_grad=False)
images, _, _ = batch
images = images.to(device)
# pdb.set_trace()
# images.size() == [1, 3, 720, 1280]
pred_target1, pred_target2 = model(images)
# pred_target1, 2 == [1, 19, 91, 161]
pred_target1 = interp_target(pred_target1)
pred_target2 = interp_target(pred_target2)
# pred_target1, 2 == [1, 19, 720, 1280]
# pdb.set_trace()
# feature_target = nn.Softmax(dim=1)(pred_target1)
# softmax_out_target = nn.Softmax(dim=1)(pred_target2)
# features = torch.cat((pred1, pred_target1), dim=0)
# outputs = torch.cat((pred2, pred_target2), dim=0)
# features.size() == [2, 19, 720, 1280]
# softmax_out.size() == [2, 19, 720, 1280]
# pdb.set_trace()
# transfer_loss = CDAN([features, softmax_out], model_D, None, None, random_layer=None)
D_out_target = CDAN(
[F.softmax(pred_target1),
F.softmax(pred_target2)],
model_D,
random_layer=None)
dc_target = torch.FloatTensor(D_out_target.size()).fill_(0).cuda()
# pdb.set_trace()
adv_loss = args.lambda_adv * nn.BCEWithLogitsLoss()(D_out_target,
dc_target)
adv_loss = adv_loss / args.iter_size
# pdb.set_trace()
# classifier_loss = nn.BCEWithLogitsLoss()(pred2,
# torch.FloatTensor(pred2.data.size()).fill_(source_label).cuda())
# pdb.set_trace()
adv_loss.backward()
adv_loss_value += adv_loss.item()
# optimizer_D.step()
#TODO: normalize loss?
for params in model_D.parameters():
params.requires_grad_(requires_grad=True)
pred1 = pred1.detach()
pred2 = pred2.detach()
D_out = CDAN([F.softmax(pred1), F.softmax(pred2)],
model_D,
random_layer=None)
dc_source = torch.FloatTensor(D_out.size()).fill_(1).cuda()
# d_loss = CDAN(D_out, dc_source, None, None, random_layer=None)
d_loss = nn.BCEWithLogitsLoss()(D_out, dc_source)
d_loss = d_loss / args.iter_size
# pdb.set_trace()
d_loss.backward()
d_loss_value += d_loss.item()
pred_target1 = pred_target1.detach()
pred_target2 = pred_target2.detach()
D_out_target = CDAN(
[F.softmax(pred_target1),
F.softmax(pred_target2)],
model_D,
random_layer=None)
dc_target = torch.FloatTensor(D_out_target.size()).fill_(0).cuda()
d_loss = nn.BCEWithLogitsLoss()(D_out_target, dc_target)
d_loss = d_loss / args.iter_size
# pdb.set_trace()
d_loss.backward()
d_loss_value += d_loss.item()
continue
optimizer.step()
optimizer_D.step()
scalar_info = {
'loss_seg1': loss_seg_value1,
'loss_seg2': loss_seg_value2,
'generator_loss': adv_loss_value,
'discriminator_loss': d_loss_value,
}
if i_iter % 10 == 0:
for key, val in scalar_info.items():
writer.add_scalar(key, val, i_iter)
# pdb.set_trace()
print('exp = {}'.format(args.snapshot_dir))
print(
'iter = {0:8d}/{1:8d}, loss_seg1 = {2:.3f} loss_seg2 = {3:.3f} generator = {4:.3f}, discriminator = {5:.3f}'
.format(i_iter, args.num_steps, loss_seg_value1, loss_seg_value2,
adv_loss_value, d_loss_value))
if i_iter >= args.num_steps_stop - 1:
print('save model ...')
torch.save(
model.state_dict(),
osp.join(args.snapshot_dir,
'GTA5_' + str(args.num_steps_stop) + '.pth'))
torch.save(
model_D.state_dict(),
osp.join(args.snapshot_dir,
'GTA5_' + str(args.num_steps_stop) + '_D.pth'))
break
if i_iter % args.save_pred_every == 0 and i_iter != 0:
print('taking snapshot ...')
torch.save(
model.state_dict(),
osp.join(args.snapshot_dir, 'GTA5_' + str(i_iter) + '.pth'))
torch.save(
model_D.state_dict(),
osp.join(args.snapshot_dir, 'GTA5_' + str(i_iter) + '_D.pth'))
check_original_discriminator(args, pred_target1, pred_target2,
i_iter)
###### also record latest saved iteration #######
args_dict['learning_rate'] = optimizer.param_groups[0]['lr']
args_dict['learning_rate_D'] = optimizer_D.param_groups[0]['lr']
args_dict['start_steps'] = i_iter
args_dict_file = args.snapshot_dir + '/args_dict_{}.json'.format(
i_iter)
with open(args_dict_file, 'w') as f:
json.dump(args_dict, f)
###### also record latest saved iteration #######
writer.close()
def drn_d_40(pretrained=False, **kwargs):
model = DRN(BasicBlock, [1, 1, 3, 4, 6, 3, 2, 2], arch='D', **kwargs)
if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['drn-d-40']))
return model
def drn_c_42(pretrained=False, **kwargs):
model = DRN(BasicBlock, [1, 1, 3, 4, 6, 3, 1, 1], arch='C', **kwargs)
if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['drn-c-42']))
return model
def main():
"""Create the model and start the evaluation process."""
args = get_arguments()
gpu0 = args.gpu
if not os.path.exists(args.save):
os.makedirs(args.save)
if args.model == 'DeeplabMulti':
model = DeeplabMulti(num_classes=args.num_classes)
elif args.model == 'DeeplabVGG':
model = DeeplabVGG(num_classes=args.num_classes)
if args.restore_from == RESTORE_FROM:
args.restore_from = RESTORE_FROM_VGG
if args.restore_from[:4] == 'http':
saved_state_dict = model_zoo.load_url(args.restore_from)
else:
saved_state_dict = torch.load(args.restore_from)
model.load_state_dict(saved_state_dict)
model.eval()
model.cuda(gpu0)
testloader = data.DataLoader(
cityscapesDataSet(args.data_dir, args.data_list, crop_size=(1024, 512), mean=IMG_MEAN, scale=False,
mirror=False, set=args.set),
batch_size=1, shuffle=False, pin_memory=True)
if version.parse(torch.__version__) >= version.parse('0.4.0'):
interp = nn.Upsample(size=(1024, 2048), mode='bilinear', align_corners=True)
else:
interp = nn.Upsample(size=(1024, 2048), mode='bilinear')
for index, batch in enumerate(testloader):
if index % 100 == 0:
print('%d processd' % index)
image, label, _, name = batch
if args.model == 'DeeplabMulti':
with torch.no_grad():
image = Variable(image)
output1, output2, _ = model(image.cuda(gpu0))
output = nn.functional.interpolate(output2, size=(1024, 2048), mode='bilinear',align_corners=True)
output = output.cpu().data[0].numpy()
elif args.model == 'DeeplabVGG':
with torch.no_grad():
image = Variable(image)
output,_ = model(image.cuda(gpu0))
output = nn.functional.interpolate(output, size=(1024, 2048), mode='bilinear', align_corners=True)
output = output.cpu().data[0].numpy()
output = output.transpose(1, 2, 0)
output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
output_col = colorize_mask(output)
output = Image.fromarray(output)
name = name[0].split('/')[-1]
output.save('%s/%s' % (args.save, name))
output_col.save('%s/%s_color.png' % (args.save, name.split('.')[0]))
def main(args):
# cudnn.enabled = True
cudnn.benchmark = True
# Setup Models
model = get_model(name=args.model, num_classes=args.num_classes)
if args.restore_from[:4] == 'http':
saved_state_dict = model_zoo.load_url(args.restore_from)
else:
saved_state_dict = torch.load(args.restore_from)
# Only copy the params that exist in current model
new_params = model.state_dict().copy()
for name, param in new_params.items():
# print(name)
if name in saved_state_dict and param.size(
) == saved_state_dict[name].size():
new_params[name].copy_(saved_state_dict[name])
# print('copy {}'.format(name))
model.load_state_dict(new_params)
# model.train()
# load model into GPU
print('GPU numbers: [ {} ] !'.format(
torch.cuda.device_count())) # GPU number
if torch.cuda.device_count() > 1:
#model = torch.nn.DataParallel(model, device_ids=range(args.num_gpus))
# model = DataParallelModel(model, device_ids=range(torch.cuda.device_count()))
model = torch.nn.DataParallel(model,
device_ids=range(
torch.cuda.device_count()))
model.cuda()
#model.to(device)
# cudnn.benchmark = True # acceleration
# Dataset and Loader
train_dataset = MITSceneParsingBenchmarkDataset(root=args.root_dataset,
split="training",
img_size=args.imgSize)
train_loader = data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=int(args.workers),
pin_memory=True,
drop_last=True)
val_dataset = MITSceneParsingBenchmarkDataset(root=args.root_dataset,
split="validation",
img_size=args.imgSize,
max_sample=args.num_val)
val_loader = data.DataLoader(
val_dataset,
batch_size=8, # args.batch_size 8
shuffle=False,
num_workers=8,
pin_memory=True,
drop_last=True)
print('1 Epoch = [ {} ] iters'.format(args.epoch_iters))
# Create loader iterator
# iterator_train = enumerate(train_loader)
#iterator_train = iter(train_loader)
crit = torch.nn.CrossEntropyLoss(ignore_index=-1)
# crit = CrossEntropyLoss(ignore_index=-1)
# Set up optimizers
optimizer = create_optimizers(model, crit, args)
if version.parse(torch.__version__) >= version.parse('0.4.0'):
interp = torch.nn.Upsample(size=(args.imgSize[0], args.imgSize[1]),
mode='bilinear',
align_corners=True)
else:
interp = torch.nn.Upsample(size=(args.imgSize[0], args.imgSize[1]),
mode='bilinear')
# interp = F.interpolate(size=(args.imgSize[0], args.imgSize[1]), mode='bilinear') # umsample
# Main loop
history = {
split: {
'epoch': [],
'loss': [],
'acc': []
}
for split in ('train', 'val')
}
# initial eval
evaluate(model, val_loader, interp, crit, history, 0, args)
for epoch in range(args.start_epoch, args.num_epoches + 1):
train(model, train_loader, interp, optimizer, crit, history, epoch,
args)
if epoch % args.eval_epoch == 0:
evaluate(model, val_loader, interp, crit, history, epoch, args)
# checkpointing
checkpoint(model, history, epoch, args)
end_time = timeit.default_timer()
print('Running time(h): [{0:.4f}]h'.format((end_time - start_time) / 3600))
print('Training Done! ***')
model.avgpool = nn.AdaptiveAvgPool2d(1)
return model
def se_resnet50(num_classes=1_000, pretrained=False):
"""Constructs a ResNet-50 model.
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
"""
model = ResNet(SEBottleneck, [3, 4, 6, 3], num_classes=num_classes)
model.avgpool = nn.AdaptiveAvgPool2d(1)
if pretrained:
model.load_state_dict(
model_zoo.load_url(
"https://www.dropbox.com/s/xpq8ne7rwa4kg4c/seresnet50-60a8950a85b2b.pkl"
))
return model
def se_resnet101(num_classes=1_000):
"""Constructs a ResNet-101 model.
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
"""
model = ResNet(SEBottleneck, [3, 4, 23, 3], num_classes=num_classes)
model.avgpool = nn.AdaptiveAvgPool2d(1)
return model
def main():
print(config)
best_mIoU = 0
if consistency_loss == 'MSE':
if len(gpus) > 1:
unlabeled_loss = torch.nn.DataParallel(MSELoss2d(),
device_ids=gpus).cuda()
else:
unlabeled_loss = MSELoss2d().cuda()
elif consistency_loss == 'CE':
if len(gpus) > 1:
unlabeled_loss = torch.nn.DataParallel(
CrossEntropyLoss2dPixelWiseWeighted(ignore_index=ignore_label),
device_ids=gpus).cuda()
else:
unlabeled_loss = CrossEntropyLoss2dPixelWiseWeighted(
ignore_index=ignore_label).cuda()
cudnn.enabled = True
# create network
model = Res_Deeplab(num_classes=num_classes)
# load pretrained parameters
#saved_state_dict = torch.load(args.restore_from)
# load pretrained parameters
if restore_from[:4] == 'http':
saved_state_dict = model_zoo.load_url(restore_from)
else:
saved_state_dict = torch.load(restore_from)
# Copy loaded parameters to model
new_params = model.state_dict().copy()
for name, param in new_params.items():
if name in saved_state_dict and param.size(
) == saved_state_dict[name].size():
new_params[name].copy_(saved_state_dict[name])
model.load_state_dict(new_params)
# init ema-model
if train_unlabeled:
ema_model = create_ema_model(model)
ema_model.train()
ema_model = ema_model.cuda()
else:
ema_model = None
if len(gpus) > 1:
if use_sync_batchnorm:
model = convert_model(model)
model = DataParallelWithCallback(model, device_ids=gpus)
else:
model = torch.nn.DataParallel(model, device_ids=gpus)
model.train()
model.cuda()
cudnn.benchmark = True
if dataset == 'cityscapes':
data_loader = get_loader('cityscapes')
data_path = get_data_path('cityscapes')
if random_crop:
data_aug = Compose([RandomCrop_city(input_size)])
else:
data_aug = None
#data_aug = Compose([RandomHorizontallyFlip()])
train_dataset = data_loader(data_path,
is_transform=True,
augmentations=data_aug,
img_size=input_size,
img_mean=IMG_MEAN)
train_dataset_size = len(train_dataset)
print('dataset size: ', train_dataset_size)
if labeled_samples is None:
trainloader = data.DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=num_workers,
pin_memory=True)
trainloader_remain = data.DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=num_workers,
pin_memory=True)
trainloader_remain_iter = iter(trainloader_remain)
else:
partial_size = labeled_samples
print('Training on number of samples:', partial_size)
np.random.seed(random_seed)
trainloader_remain = data.DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=4,
pin_memory=True)
trainloader_remain_iter = iter(trainloader_remain)
#New loader for Domain transfer
if True:
data_loader = get_loader('gta')
data_path = get_data_path('gta')
if random_crop:
data_aug = Compose([RandomCrop_gta(input_size)])
else:
data_aug = None
#data_aug = Compose([RandomHorizontallyFlip()])
train_dataset = data_loader(data_path,
list_path='./data/gta5_list/train.txt',
augmentations=data_aug,
img_size=(1280, 720),
mean=IMG_MEAN)
trainloader = data.DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=num_workers,
pin_memory=True)
trainloader_iter = iter(trainloader)
print('gta size:', len(trainloader))
#Load new data for domain_transfer
# optimizer for segmentation network
learning_rate_object = Learning_Rate_Object(
config['training']['learning_rate'])
if optimizer_type == 'SGD':
if len(gpus) > 1:
optimizer = optim.SGD(
model.module.optim_parameters(learning_rate_object),
lr=learning_rate,
momentum=momentum,
weight_decay=weight_decay)
else:
optimizer = optim.SGD(model.optim_parameters(learning_rate_object),
lr=learning_rate,
momentum=momentum,
weight_decay=weight_decay)
elif optimizer_type == 'Adam':
if len(gpus) > 1:
optimizer = optim.Adam(
model.module.optim_parameters(learning_rate_object),
lr=learning_rate,
momentum=momentum,
weight_decay=weight_decay)
else:
optimizer = optim.Adam(
model.optim_parameters(learning_rate_object),
lr=learning_rate,
weight_decay=weight_decay)
optimizer.zero_grad()
interp = nn.Upsample(size=(input_size[0], input_size[1]),
mode='bilinear',
align_corners=True)
start_iteration = 0
if args.resume:
start_iteration, model, optimizer, ema_model = _resume_checkpoint(
args.resume, model, optimizer, ema_model)
accumulated_loss_l = []
accumulated_loss_u = []
if not os.path.exists(checkpoint_dir):
os.makedirs(checkpoint_dir)
with open(checkpoint_dir + '/config.json', 'w') as handle:
json.dump(config, handle, indent=4, sort_keys=True)
epochs_since_start = 0
for i_iter in range(start_iteration, num_iterations):
model.train()
loss_u_value = 0
loss_l_value = 0
optimizer.zero_grad()
if lr_schedule:
adjust_learning_rate(optimizer, i_iter)
# training loss for labeled data only
try:
batch = next(trainloader_iter)
if batch[0].shape[0] != batch_size:
batch = next(trainloader_iter)
except:
epochs_since_start = epochs_since_start + 1
print('Epochs since start: ', epochs_since_start)
trainloader_iter = iter(trainloader)
batch = next(trainloader_iter)
#if random_flip:
# weak_parameters={"flip":random.randint(0,1)}
#else:
weak_parameters = {"flip": 0}
images, labels, _, _ = batch
images = images.cuda()
labels = labels.cuda().long()
#images, labels = weakTransform(weak_parameters, data = images, target = labels)
pred = interp(model(images))
L_l = loss_calc(pred, labels) # Cross entropy loss for labeled data
#L_l = torch.Tensor([0.0]).cuda()
if train_unlabeled:
try:
batch_remain = next(trainloader_remain_iter)
if batch_remain[0].shape[0] != batch_size:
batch_remain = next(trainloader_remain_iter)
except:
trainloader_remain_iter = iter(trainloader_remain)
batch_remain = next(trainloader_remain_iter)
images_remain, _, _, _, _ = batch_remain
images_remain = images_remain.cuda()
inputs_u_w, _ = weakTransform(weak_parameters, data=images_remain)
#inputs_u_w = inputs_u_w.clone()
logits_u_w = interp(ema_model(inputs_u_w))
logits_u_w, _ = weakTransform(
getWeakInverseTransformParameters(weak_parameters),
data=logits_u_w.detach())
pseudo_label = torch.softmax(logits_u_w.detach(), dim=1)
max_probs, targets_u_w = torch.max(pseudo_label, dim=1)
if mix_mask == "class":
for image_i in range(batch_size):
classes = torch.unique(labels[image_i])
#classes=classes[classes!=ignore_label]
nclasses = classes.shape[0]
#if nclasses > 0:
classes = (classes[torch.Tensor(
np.random.choice(nclasses,
int((nclasses + nclasses % 2) / 2),
replace=False)).long()]).cuda()
if image_i == 0:
MixMask0 = transformmasks.generate_class_mask(
labels[image_i], classes).unsqueeze(0).cuda()
else:
MixMask1 = transformmasks.generate_class_mask(
labels[image_i], classes).unsqueeze(0).cuda()
elif mix_mask == None:
MixMask = torch.ones((inputs_u_w.shape))
strong_parameters = {"Mix": MixMask0}
if random_flip:
strong_parameters["flip"] = random.randint(0, 1)
else:
strong_parameters["flip"] = 0
if color_jitter:
strong_parameters["ColorJitter"] = random.uniform(0, 1)
else:
strong_parameters["ColorJitter"] = 0
if gaussian_blur:
strong_parameters["GaussianBlur"] = random.uniform(0, 1)
else:
strong_parameters["GaussianBlur"] = 0
inputs_u_s0, _ = strongTransform(
strong_parameters,
data=torch.cat(
(images[0].unsqueeze(0), images_remain[0].unsqueeze(0))))
strong_parameters["Mix"] = MixMask1
inputs_u_s1, _ = strongTransform(
strong_parameters,
data=torch.cat(
(images[1].unsqueeze(0), images_remain[1].unsqueeze(0))))
inputs_u_s = torch.cat((inputs_u_s0, inputs_u_s1))
logits_u_s = interp(model(inputs_u_s))
strong_parameters["Mix"] = MixMask0
_, targets_u0 = strongTransform(strong_parameters,
target=torch.cat(
(labels[0].unsqueeze(0),
targets_u_w[0].unsqueeze(0))))
strong_parameters["Mix"] = MixMask1
_, targets_u1 = strongTransform(strong_parameters,
target=torch.cat(
(labels[1].unsqueeze(0),
targets_u_w[1].unsqueeze(0))))
targets_u = torch.cat((targets_u0, targets_u1)).long()
if pixel_weight == "threshold_uniform":
unlabeled_weight = torch.sum(
max_probs.ge(0.968).long() == 1).item() / np.size(
np.array(targets_u.cpu()))
pixelWiseWeight = unlabeled_weight * torch.ones(
max_probs.shape).cuda()
elif pixel_weight == "threshold":
pixelWiseWeight = max_probs.ge(0.968).float().cuda()
elif pixel_weight == False:
pixelWiseWeight = torch.ones(max_probs.shape).cuda()
onesWeights = torch.ones((pixelWiseWeight.shape)).cuda()
strong_parameters["Mix"] = MixMask0
_, pixelWiseWeight0 = strongTransform(
strong_parameters,
target=torch.cat((onesWeights[0].unsqueeze(0),
pixelWiseWeight[0].unsqueeze(0))))
strong_parameters["Mix"] = MixMask1
_, pixelWiseWeight1 = strongTransform(
strong_parameters,
target=torch.cat((onesWeights[1].unsqueeze(0),
pixelWiseWeight[1].unsqueeze(0))))
pixelWiseWeight = torch.cat(
(pixelWiseWeight0, pixelWiseWeight1)).cuda()
if consistency_loss == 'MSE':
unlabeled_weight = torch.sum(
max_probs.ge(0.968).long() == 1).item() / np.size(
np.array(targets_u.cpu()))
#pseudo_label = torch.cat((pseudo_label[1].unsqueeze(0),pseudo_label[0].unsqueeze(0)))
L_u = consistency_weight * unlabeled_weight * unlabeled_loss(
logits_u_s, pseudo_label)
elif consistency_loss == 'CE':
L_u = consistency_weight * unlabeled_loss(
logits_u_s, targets_u, pixelWiseWeight)
loss = L_l + L_u
else:
loss = L_l
if len(gpus) > 1:
#print('before mean = ',loss)
loss = loss.mean()
#print('after mean = ',loss)
loss_l_value += L_l.mean().item()
if train_unlabeled:
loss_u_value += L_u.mean().item()
else:
loss_l_value += L_l.item()
if train_unlabeled:
loss_u_value += L_u.item()
loss.backward()
optimizer.step()
# update Mean teacher network
if ema_model is not None:
alpha_teacher = 0.99
ema_model = update_ema_variables(ema_model=ema_model,
model=model,
alpha_teacher=alpha_teacher,
iteration=i_iter)
print(
'iter = {0:6d}/{1:6d}, loss_l = {2:.3f}, loss_u = {3:.3f}'.format(
i_iter, num_iterations, loss_l_value, loss_u_value))
if i_iter % save_checkpoint_every == 0 and i_iter != 0:
if epochs_since_start * len(trainloader) < save_checkpoint_every:
_save_checkpoint(i_iter,
model,
optimizer,
config,
ema_model,
overwrite=False)
else:
_save_checkpoint(i_iter, model, optimizer, config, ema_model)
if config['utils']['tensorboard']:
if 'tensorboard_writer' not in locals():
tensorboard_writer = tensorboard.SummaryWriter(log_dir,
flush_secs=30)
accumulated_loss_l.append(loss_l_value)
if train_unlabeled:
accumulated_loss_u.append(loss_u_value)
if i_iter % log_per_iter == 0 and i_iter != 0:
tensorboard_writer.add_scalar('Training/Supervised loss',
np.mean(accumulated_loss_l),
i_iter)
accumulated_loss_l = []
if train_unlabeled:
tensorboard_writer.add_scalar('Training/Unsupervised loss',
np.mean(accumulated_loss_u),
i_iter)
accumulated_loss_u = []
if i_iter % val_per_iter == 0 and i_iter != 0:
model.eval()
if dataset == 'cityscapes':
mIoU, eval_loss = evaluate(model,
dataset,
ignore_label=250,
input_size=(512, 1024),
save_dir=checkpoint_dir)
model.train()
if mIoU > best_mIoU and save_best_model:
best_mIoU = mIoU
_save_checkpoint(i_iter,
model,
optimizer,
config,
ema_model,
save_best=True)
if config['utils']['tensorboard']:
tensorboard_writer.add_scalar('Validation/mIoU', mIoU, i_iter)
tensorboard_writer.add_scalar('Validation/Loss', eval_loss,
i_iter)
if save_unlabeled_images and train_unlabeled and i_iter % save_checkpoint_every == 0:
# Saves two mixed images and the corresponding prediction
save_image(inputs_u_s[0].cpu(), i_iter, 'input1',
palette.CityScpates_palette)
save_image(inputs_u_s[1].cpu(), i_iter, 'input2',
palette.CityScpates_palette)
_, pred_u_s = torch.max(logits_u_s, dim=1)
save_image(pred_u_s[0].cpu(), i_iter, 'pred1',
palette.CityScpates_palette)
save_image(pred_u_s[1].cpu(), i_iter, 'pred2',
palette.CityScpates_palette)
_save_checkpoint(num_iterations, model, optimizer, config, ema_model)
model.eval()
if dataset == 'cityscapes':
mIoU, val_loss = evaluate(model,
dataset,
ignore_label=250,
input_size=(512, 1024),
save_dir=checkpoint_dir)
model.train()
if mIoU > best_mIoU and save_best_model:
best_mIoU = mIoU
_save_checkpoint(i_iter,
model,
optimizer,
config,
ema_model,
save_best=True)
if config['utils']['tensorboard']:
tensorboard_writer.add_scalar('Validation/mIoU', mIoU, i_iter)
tensorboard_writer.add_scalar('Validation/Loss', val_loss, i_iter)
end = timeit.default_timer()
print('Total time: ' + str(end - start) + 'seconds')
def build_resnet_backbone(cfg):
"""
Create a ResNet instance from config.
Returns:
ResNet: a :class:`ResNet` instance.
"""
# fmt: off
pretrain = cfg.MODEL.BACKBONE.PRETRAIN
pretrain_path = cfg.MODEL.BACKBONE.PRETRAIN_PATH
last_stride = cfg.MODEL.BACKBONE.LAST_STRIDE
bn_norm = cfg.MODEL.BACKBONE.NORM
num_splits = cfg.MODEL.BACKBONE.NORM_SPLIT
with_ibn = cfg.MODEL.BACKBONE.WITH_IBN
with_se = cfg.MODEL.BACKBONE.WITH_SE
with_nl = cfg.MODEL.BACKBONE.WITH_NL
depth = cfg.MODEL.BACKBONE.DEPTH
num_blocks_per_stage = {
34: [3, 4, 6, 3],
50: [3, 4, 6, 3],
101: [3, 4, 23, 3],
152: [3, 8, 36, 3],
}[depth]
nl_layers_per_stage = {
34: [3, 4, 6, 3],
50: [0, 2, 3, 0],
101: [0, 2, 9, 0]
}[depth]
block = {34: BasicBlock, 50: Bottleneck, 101: Bottleneck}[depth]
model = ResNet(last_stride, bn_norm, num_splits, with_ibn, with_se,
with_nl, block, num_blocks_per_stage, nl_layers_per_stage)
if pretrain:
if not with_ibn:
try:
state_dict = torch.load(
pretrain_path, map_location=torch.device('cpu'))['model']
# Remove module.encoder in name
new_state_dict = {}
for k in state_dict:
new_k = '.'.join(k.split('.')[2:])
if new_k in model.state_dict() and (
model.state_dict()[new_k].shape
== state_dict[k].shape):
new_state_dict[new_k] = state_dict[k]
state_dict = new_state_dict
logger.info(f"Loading pretrained model from {pretrain_path}")
except FileNotFoundError or KeyError:
# original resnet
state_dict = model_zoo.load_url(model_urls[depth])
logger.info("Loading pretrained model from torchvision")
else:
state_dict = torch.load(
pretrain_path,
map_location=torch.device('cpu'))['state_dict'] # ibn-net
# Remove module in name
new_state_dict = {}
for k in state_dict:
new_k = '.'.join(k.split('.')[1:])
if new_k in model.state_dict() and (
model.state_dict()[new_k].shape
== state_dict[k].shape):
new_state_dict[new_k] = state_dict[k]
state_dict = new_state_dict
logger.info(f"Loading pretrained model from {pretrain_path}")
incompatible = model.load_state_dict(state_dict, strict=False)
if incompatible.missing_keys:
logger.info(
get_missing_parameters_message(incompatible.missing_keys))
if incompatible.unexpected_keys:
logger.info(
get_unexpected_parameters_message(
incompatible.unexpected_keys))
return model
def init_network(model='resnet101',
pooling='gem',
whitening=False,
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225],
pretrained=True):
# loading network from torchvision
if pretrained:
if model not in FEATURES:
# initialize with network pretrained on imagenet in pytorch
net_in = getattr(torchvision.models, model)(pretrained=True)
else:
# initialize with random weights, later on we will fill features with custom pretrained network
net_in = getattr(torchvision.models, model)(pretrained=False)
else:
# initialize with random weights
net_in = getattr(torchvision.models, model)(pretrained=False)
# initialize features
# take only convolutions for features,
# always ends with ReLU to make last activations non-negative
if model.startswith('alexnet'):
features = list(net_in.features.children())[:-1]
elif model.startswith('vgg'):
features = list(net_in.features.children())[:-1]
elif model.startswith('resnet'):
features = list(net_in.children())[:-2]
elif model.startswith('densenet'):
features = list(net_in.features.children())
features.append(nn.ReLU(inplace=True))
elif model.startswith('squeezenet'):
features = list(net_in.features.children())
else:
raise ValueError('Unsupported or unknown model: {}!'.format(model))
# initialize pooling
pool = POOLING[pooling]()
# get output dimensionality size
dim = OUTPUT_DIM[model]
# initialize whitening
if whitening:
w = '{}-{}'.format(model, pooling)
whiten = nn.Linear(dim, dim, bias=True)
if w in WHITENING:
print(">> {}: for '{}' custom computed whitening '{}' is used".
format(os.path.basename(__file__), w,
os.path.basename(WHITENING[w])))
whiten_dir = os.path.join(get_data_root(), 'whiten')
whiten.load_state_dict(
model_zoo.load_url(WHITENING[w], model_dir=whiten_dir))
else:
print(
">> {}: for '{}' there is no whitening computed, random weights are used"
.format(os.path.basename(__file__), w))
else:
whiten = None
# create meta information to be stored in the network
meta = {
'architecture': model,
'pooling': pooling,
'whitening': whitening,
'outputdim': dim,
'mean': mean,
'std': std
}
# create a generic image retrieval network
net = ImageRetrievalNet(features, pool, whiten, meta)
# initialize features with custom pretrained network if needed
if pretrained and model in FEATURES:
print(
">> {}: for '{}' custom pretrained features '{}' are used".format(
os.path.basename(__file__), model,
os.path.basename(FEATURES[model])))
model_dir = os.path.join(get_data_root(), 'networks')
net.features.load_state_dict(
model_zoo.load_url(FEATURES[model], model_dir=model_dir))
return net
def resnet18(pretrained=False, **kwargs):
model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs)
if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['resnet18']))
return model
def resnet152(pretrained=True):
model = ResNet(Bottleneck, [3, 8, 36, 3])
if pretrained:
load_weights_sequential(model,
model_zoo.load_url(model_urls['resnet152']))
return model
