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