def __init__(self, checkpoint_path='checkpoint/checkpoint.pth', imsize=512): super().__init__() self.conf_threshold = 0.6 # Any pixel below this confidence value will not be considered as a detection. self.checkpoint_path = checkpoint_path self.transform = transforms.Compose( [transforms.Resize((imsize, imsize)), transforms.ToTensor()]) # Create Model self.numInputChannels = 3 self.num_of_classes = 2 self.output_stride = 8 # Chosen while training the model. Doesn't make diff in inference. Possible Values: 8, 16 self.net = deeplab_xception.DeepLabv3_plus( nInputChannels=self.numInputChannels, n_classes=self.num_of_classes, os=self.output_stride, pretrained=False) self.net = nn.DataParallel(self.net) self.net.load_state_dict(torch.load(self.checkpoint_path)) # Send model to Device (GPU or CPU) self.device = torch.device( "cuda:0" if torch.cuda.is_available() else "cpu") self.net.to(self.device) self.net.eval()
save_dir_root = os.path.join(os.path.dirname(os.path.abspath(__file__))) exp_name = os.path.dirname(os.path.abspath(__file__)).split('/')[-1] if resume_epoch != 0: runs = sorted(glob.glob(os.path.join(save_dir_root, 'run', 'run_*'))) run_id = int(runs[-1].split('_')[-1]) if runs else 0 else: runs = sorted(glob.glob(os.path.join(save_dir_root, 'run', 'run_*'))) run_id = int(runs[-1].split('_')[-1]) + 1 if runs else 0 save_dir = os.path.join(save_dir_root, 'run', 'run_' + str(run_id)) # Network definition net = deeplab_xception.DeepLabv3_plus(nInputChannels=3, n_classes=21, pretrained=True) modelName = 'deeplabv3+' criterion = utils.cross_entropy2d if resume_epoch == 0: print("Training deeplabv3+ from scratch...") else: print("Initializing weights from: {}...".format( os.path.join(save_dir, 'models', modelName + '_epoch-' + str(resume_epoch - 1) + '.pth'))) net.load_state_dict( torch.load(os.path.join( save_dir, 'models', modelName + '_epoch-' + str(resume_epoch - 1) + '.pth'), map_location=lambda storage, loc: storage)
p['lr'] = 1e-10 # Learning rate p['wd'] = 5e-4 # Weight decay p['momentum'] = 0.9 # Momentum p['epoch_size'] = 2 # How many epochs to change learning rate p['Model'] = 'deeplab' # Choose model: unet or deeplab backbone = 'xception' # For deeplab only: Use xception or resnet as feature extractor, num_of_classes = 2 imsize = 512 # 256 or 512 output_stride = 8 # 8 or 16, 8 is better. Controls output stride of the deeplab model, which increases resolution of convolutions. numInputChannels = 3 # Network definition if p['Model'] == 'deeplab': if backbone == 'xception': net = deeplab_xception.DeepLabv3_plus(nInputChannels=numInputChannels, n_classes=num_of_classes, os=output_stride, pretrained=True) elif backbone == 'resnet': net = deeplab_resnet.DeepLabv3_plus(nInputChannels=numInputChannels, n_classes=num_of_classes, os=output_stride, pretrained=True) else: raise NotImplementedError modelName = 'deeplabv3plus-' + backbone # Use the following optimizer optimizer = optim.SGD(net.parameters(), lr=p['lr'], momentum=p['momentum'], weight_decay=p['wd']) p['optimizer'] = str(optimizer) # Use the following loss function criterion = utils.cross_entropy2d else: raise NotImplementedError
save_dir_root = os.path.join(os.path.dirname(os.path.abspath(__file__))) exp_name = os.path.dirname(os.path.abspath(__file__)).split('/')[-1] if resume_epoch != 0: runs = sorted(glob.glob(os.path.join(save_dir_root, 'run', 'run_*'))) run_id = int(runs[-1].split('_')[-1]) if runs else 0 else: runs = sorted(glob.glob(os.path.join(save_dir_root, 'run', 'run_*'))) run_id = int(runs[-1].split('_')[-1]) + 1 if runs else 0 save_dir = os.path.join(save_dir_root, 'run') # Network definition if backbone == 'xception': net = deeplab_xception.DeepLabv3_plus(nInputChannels=3, n_classes=12, os=16, pretrained=False) elif backbone == 'resnet': net = deeplab_resnet.DeepLabv3_plus(nInputChannels=3, n_classes=12, os=16, pretrained=False) else: raise NotImplementedError modelName = 'deeplabv3plus-' + backbone + '-voc' criterion = loss.dice_bce_loss() # criterion = loss.FocalLoss(gamma=2) print("Backbone network is {}".format(backbone))
p['momentum'] = 0.9 # Momentum p['epoch_size'] = 40 # How many epochs to change learning rate save_dir_root = os.path.join(os.path.dirname(os.path.abspath(__file__))) exp_name = os.path.dirname(os.path.abspath(__file__)).split('/')[-1] if resume_epoch == 0: runs = sorted(glob.glob(os.path.join(save_dir_root, 'run', 'run_*'))) run_id = int(runs[-1].split('_')[-1]) + 1 if runs else 0 else: run_id = 0 save_dir = os.path.join(save_dir_root, 'run', 'run_' + str(run_id)) # Network definition net = deeplab_xception.DeepLabv3_plus(nInputChannels=3, n_classes=21) modelName = 'deeplabv3+' criterion = utils.cross_entropy2d if resume_epoch == 0: print("Training deeplabv3+ from scratch...") else: print("Initializing weights from: {}...".format( os.path.join(save_dir, 'models', modelName + '_epoch-' + str(resume_epoch - 1) + '.pth'))) net.load_state_dict( torch.load(os.path.join(save_dir, 'models', modelName + '_epoch-' + str(resume_epoch - 1) + '.pth'), map_location=lambda storage, loc: storage)) # Load all tensors onto the CPU if gpu_id >= 0: torch.cuda.set_device(device=gpu_id) net.cuda()