def test_fcn(self): x = Variable(torch.randn(BATCH_SIZE, 3, 224, 224).fill_(1.0)) self.exportTest(toC( fcn_resnet101(pretrained=False, pretrained_backbone=False)), toC(x), rtol=1e-3, atol=1e-5)
def get_model(model_name: str, device, num_classes: int, pretrained: bool = False) -> torch.nn: """ Gets the model requested from the config file :param model_name: name of model used in the program :param device: name of device model is hosted on :param num_classes: number of classes in the model (including background 0,0,0) :param pretrained: boolean of starting on a pretrained model :return: the nn.model to be trained """ print(pretrained) model_dict = {'fcn_resnet50': models.fcn_resnet50(pretrained=pretrained, num_classes=num_classes).to(device).eval(), 'fcn_resnet101': models.fcn_resnet101(pretrained=pretrained, num_classes=num_classes).to(device).eval(), 'deeplabv3_resnet50': models.deeplabv3_resnet50(pretrained=pretrained, num_classes=num_classes).to(device).eval(), 'deeplabv3_resnet101': models.deeplabv3_resnet101(pretrained=pretrained, num_classes=num_classes).to(device).eval(), 'deeplabv3_mobilenet_v3_large': models.deeplabv3_mobilenet_v3_large(pretrained=pretrained, num_classes=num_classes).to( device).eval(), 'lraspp_mobilenet_v3_large': models.lraspp_mobilenet_v3_large(pretrained=pretrained, num_classes=num_classes).to( device).eval() } try: model = model_dict[model_name] except KeyError: print(f"KeyError, model_name is not valid allowable names are: {model_dict.keys()}") model = None model = model.eval() return model
def __init__(self): super(Encoder101, self).__init__() self.resnet = deeplabv3_resnet101(pretrained=True) self.conv1_p = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=True) resnet = fcn_resnet101(pretrained=True) self.conv1 = resnet.backbone.conv1 self.bn1 = resnet.backbone.bn1 self.relu = resnet.backbone.relu # 1/2, 64 self.maxpool = resnet.backbone.maxpool self.res2 = resnet.backbone.layer1 # 1/4, 256 self.res3 = resnet.backbone.layer2 # 1/8, 512 self.res4 = resnet.backbone.layer3 # 1/16, 1024 self.res5 = resnet.backbone.layer4 # 1/32, 2048 self.register_buffer( 'mean', torch.FloatTensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)) self.register_buffer( 'std', torch.FloatTensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1))
def load_seg_model(cfg, load_model_only=False): modeltype = cfg['training']['model'] if modeltype == 'fcn_resnet101': n_channels = len(cfg['data']['input_classes']) n_outputs = len(cfg['data']['output_classes']) model = fcn_resnet101(pretrained=False, num_classes=n_outputs) model.backbone.conv1 = nn.Conv2d(n_channels, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False) elif modeltype == 'deeplabv3_resnet101': n_channels = len(cfg['data']['input_classes']) n_outputs = len(cfg['data']['output_classes']) aux_loss = cfg['training']['aux_loss'] is not False model = deeplabv3_resnet101(pretrained=False, num_classes=n_outputs, aux_loss=aux_loss) model.backbone.conv1 = nn.Conv2d(n_channels, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False) elif modeltype == 'higher_hrnet': hrnet_cfg = get_hrnet_cfg(cfg) model = get_hrnet_model(hrnet_cfg) else: raise ValueError() if modelpath := cfg['resume'].get('path', None): state = torch.load(modelpath) model.load_state_dict(state['state_dict']) if load_model_only: return model starting_epoch = state['epoch'] if conf_epoch := cfg['resume'].get('epoch', None): print( f"WARNING: Loaded model trained for {starting_epoch - 1} epochs but config explicitly overrides to {conf_epoch}") starting_epoch = conf_epoch
def resume(self, file, test=False): import torch if test and not file: self.fcn = fcn_resnet101(pretrained=True, num_classes=2) return if file: print('Loading checkpoint from: ' + file) checkpoint = torch.load(file) checkpoint = checkpoint['model_state_dict'] self.load_state_dict(checkpoint)
def __init__(self, n_channels=3, n_classes=21, softmax_out=True, resnet_type=101, pretrained=False): super(FCN, self).__init__() self.resnet_type = resnet_type self.n_channels = n_channels self.n_classes = n_classes self.pretrained = pretrained # Input conv is applied to convert the input to 3 ch depth self.inconv = None if n_channels != 3: self.inconv = FwdConv(n_channels, 3, kernel_size=1, padding=0) # Pre-trained model needs to be an identical network if pretrained: mid_classes = 21 else: mid_classes = n_classes # Maind body if resnet_type == 50: self.fcn_body = fcn_resnet50(pretrained=False, num_classes=mid_classes) self.pretrained = False else: self.fcn_body = fcn_resnet101(pretrained=pretrained, num_classes=mid_classes) if n_classes != 21: self.fcn_body.classifier[-1] = nn.Conv2d(512, n_classes, kernel_size=(1, 1), stride=(1, 1)) if self.fcn_body.aux_classifier != None: self.fcn_body.aux_classifier[-1] = nn.Conv2d(512, n_classes, kernel_size=(1, 1), stride=(1, 1)) # Softmax alternative self.has_softmax = softmax_out if softmax_out: self.softmax = nn.Softmax2d() else: self.softmax = None
def __init__(self, end_module, keyword='out', pretrain=False, freeze=False): super(ResNet101_FCN, self).__init__() self.resnet = models.fcn_resnet101(pretrained=pretrain, progress=True, num_classes=21) if freeze: for module in self.resnet.parameters(): module.requires_grad = False print("module frozen.") self.end_module = end_module self.output = None self.key = keyword
def __init__(self): self.model = seg.fcn_resnet101(pretrained=True, progress=True) self.model.cuda() self.model.eval()
def __init__(self, n_class=21): super(FCN, self).__init__() self.fcn = fcn_resnet101(pretrained=False, num_classes=n_class)
def fcn_resnet101(num_classes): """Return torchvision fcn_resnet101 and wrap it in SingleOutNet.""" return SingleOutNet( seg_models.fcn_resnet101(num_classes=num_classes, aux_loss=True))
def train(opt): # Specify the Model Architecture if opt.architecture.lower() == "deeplabv3_resnet50": model = models.deeplabv3_resnet50(pretrained=True, progress=True) model.classifier = models.deeplabv3.DeepLabHead(2048, 4) elif opt.architecture.lower() == "deeplabv3_resnet101": model = models.deeplabv3_resnet101(pretrained=True, progress=True) model.classifier = models.deeplabv3.DeepLabHead(2048, 4) elif opt.architecture.lower() == "fcn_resnet50": model = models.fcn_resnet50(pretrained=True, progress=True) model.classifier = models.fcn.FCNHead(2048, 4) elif opt.architecture.lower() == "fcn_resnet101": model = models.fcn_resnet101(pretrained=True, progress=True) model.classifier = models.fcn.FCNHead(2048, 4) # Define Optimizer if opt.optim.lower() == "adam": modelOptim = torch.optim.Adam(model.parameters(), lr=opt.lr) elif opt.optim.lower() == "sgd": modelOptim = torch.optim.SGD(model.parameters(), lr=opt.lr) # Define Loss Function lossFnc = torch.nn.CrossEntropyLoss() # Set Training and Validation Datasets dataTransforms = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) segdata = { x: Dataset(root=Path(opt.data_path) / x, imageFolder="images", maskFolder="masks", transforms=dataTransforms) for x in ["train", "valid"] } dataLoaders = { x: DataLoader(segdata[x], batch_size=opt.batch_size, shuffle=True, num_workers=opt.num_workers) for x in ["train", "valid"] } # Set Training Device device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") model.to(device) # Create and Initialize Training Log File #perfMetrics = {"f1-score": f1_score, "auroc": roc_auc_score} perfMetrics = {"f1-score": f1_score} fieldnames = ['epoch', 'train_loss', 'valid_loss'] + \ [f'train_{m}' for m in perfMetrics.keys()] + \ [f'valid_{m}' for m in perfMetrics.keys()] with open(os.path.join('log.csv'), 'w', newline='') as csvfile: writer = csv.DictWriter(csvfile, fieldnames=fieldnames) writer.writeheader() #model = torch.load('trainedModel25.pth') # Train startTimer = time.time() for epoch in range(1, opt.num_epochs + 1): print('-' * 60) print("Epoch: {}/{}".format(epoch, opt.num_epochs)) batchsummary = {a: [0] for a in fieldnames} for phase in ["train", "valid"]: if phase == 'train': model.train() else: model.eval() # Iterate over data. for sample in tqdm(iter(dataLoaders[phase]), file=sys.stdout): inputs = sample['image'].to(device) masks = sample['mask'].to(device) # zero the parameter gradients modelOptim.zero_grad() # track history if only in train with torch.set_grad_enabled(phase == 'train'): outputs = model(inputs) loss = lossFnc(outputs['out'], masks) y_pred = outputs['out'].data.cpu().numpy().ravel() y_true = masks.data.cpu().numpy().ravel() for name, metric in perfMetrics.items(): if name == 'f1-score': # Use a classification threshold of 0.1 f1Classes = np.zeros(4) nPixels = np.zeros(4) for classID in range(4): f1Classes[classID] = metric( y_true == classID, y_pred[classID * len(y_true): (classID + 1) * len(y_true)] > 0.1) nPixels[classID] = np.count_nonzero( y_true == classID) f1weights = nPixels / (np.sum(nPixels)) f1 = np.matmul(f1Classes, f1weights) batchsummary[f'{phase}_{name}'].append(f1) else: batchsummary[f'{phase}_{name}'].append( metric(y_true.astype('uint8'), y_pred)) # backward + optimize only if in training phase if phase == 'train': loss.backward() modelOptim.step() batchsummary['epoch'] = epoch epoch_loss = loss batchsummary[f'{phase}_loss'] = epoch_loss.item() for field in fieldnames[3:]: batchsummary[field] = np.mean(batchsummary[field]) print((f'train loss: {batchsummary["train_loss"]: .4f}, ' f'valid loss: {batchsummary["valid_loss"]: .4f}, ' f'train f1-score: {batchsummary["train_f1-score"]: .4f}, ' f'valid f1-score: {batchsummary["valid_f1-score"]: .4f}, ')) #f'train auroc: {batchsummary["train_auroc"]: .4f}, ' #f'valid auroc: {batchsummary["valid_auroc"]: .4f}, ')) with open(os.path.join('log.csv'), 'a', newline='') as csvfile: writer = csv.DictWriter(csvfile, fieldnames=fieldnames) writer.writerow(batchsummary) # deep copy the model if phase == 'valid' and loss < 1e10: best_loss = loss best_model_wts = copy.deepcopy(model.state_dict()) time_elapsed = time.time() - startTimer print('-' * 60) print('Training completed in {:.0f}m {:.0f}s'.format( time_elapsed // 60, time_elapsed % 60)) print(f'Lowest validation loss: {best_loss: .4f}') # load best model weights model.load_state_dict(best_model_wts)
def test_fcn(self): x = Variable(torch.randn(BATCH_SIZE, 3, 224, 224).fill_(1.0)) self.exportTest(toC(fcn_resnet101()), toC(x), rtol=1e-3, atol=1e-5)
def predict(config_file, model_file, output_dir, n_classes, architecture, multi_crop=False, aux_loss=None): pred_dir = Path(output_dir) / 'predictions' col_dir = Path(output_dir) / 'colorized' pred_dir.mkdir(exist_ok=True) col_dir.mkdir(exist_ok=True) if architecture == 'deeplab': model = deeplabv3_resnet101(pretrained=False, progress=True, num_classes=n_classes, aux_loss=aux_loss) elif architecture == 'fcn': model = fcn_resnet101(pretrained=False, progress=True, num_classes=n_classes, aux_loss=aux_loss) else: raise ValueError('Unknown architecture specified. ' 'Please choose either \'fcn\' or \'deeplab\'.') model.load_state_dict(torch.load(model_file)) model.eval() model.cuda() criterion = CrossEntropyLoss(ignore_index=255) test_loader, data_list = get_val_loader(config_file) data_loader = tqdm(test_loader, total=len(test_loader), ncols=80) color_palette = np.loadtxt( 'semseg/data/cityscapes/cityscapes_colors.txt').astype('uint8') if not multi_crop: for i, (img, label) in enumerate(data_loader): img_path, label_path = data_list[i] img = torch.cuda.FloatTensor(img.numpy()) label = torch.cuda.FloatTensor(label.numpy()).long() output = model(img)['out'] loss = criterion(output, label) data_loader.set_postfix(loss=loss.item()) pred_path = pred_dir / Path(img_path).name col_path = col_dir / Path(img_path).name # Apply argmax to prediction image. The resulting indices # represent an 8-bit image with each pixel value corresponding # to the respective class label values, indices = torch.max(torch.squeeze(output), 0) pred = indices.cpu().numpy() color_pred = colorize(pred, color_palette) cv2.imwrite(str(pred_path), pred) color_pred.save(str(col_path)) else: test_args = config.load_cfg_from_cfg_file(config_file) mean, std = get_mean_std(test_args.mean_std) crop_h = test_args.train_h crop_w = test_args.train_w logger = get_logger() predict_multi_crop(logger, test_loader, data_list, model, n_classes, mean, test_args.base_size, crop_h, crop_w, test_args.scales, pred_dir, col_dir, color_palette)
type=str, default='cityscapes_deeplab.yaml', help='config file') args0 = parser.parse_args() cudnn.benchmark = True args, data_loader = get_train_loader(args0.config) if args.architecture == 'deeplab': model = deeplabv3_resnet101(pretrained=False, progress=True, num_classes=args.classes, aux_loss=aux_loss) elif args.architecture == 'fcn': model = fcn_resnet101(pretrained=False, progress=True, num_classes=args.classes, aux_loss=aux_loss) else: raise ValueError('Unknown architecture specified. ' 'Please choose either \'fcn\' or \'deeplab\'.') model.cuda() print(model) criterion = CrossEntropyLoss(ignore_index=args.ignore_label) optimizer = optim.SGD(model.parameters(), lr=args.base_lr, momentum=args.momentum) loss_log = [] for epoch in range(args.epochs):
def __init__(self): super(FCNResNet101Model, self).__init__() self.model = fcn_resnet101(pretrained=True, progress=False) self.MEAN = torch.tensor([0.485, 0.456, 0.406]) self.STD = torch.tensor([0.229, 0.224, 0.225])
#%% import torch from torchvision.models import segmentation from PIL import Image import numpy as np from matplotlib import pyplot as plt #%% device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') #%% fcn_res101 = segmentation.fcn_resnet101(pretrained=True) fcn_res101.to(device) #%% print(fcn_res101) # %% img = Image.open("1.jpg") img = np.array(img) img_tensor = torch.from_numpy(img.astype(np.float32)) img_tensor = torch.unsqueeze(img_tensor,0) img_tensor = img_tensor.permute(0,3,1,2) img_tensor = img_tensor.to(device) #%% predict = fcn_res101(img_tensor) #%% pre_numpy = predict['out'].detach().numpy() pre_numpy_person = pre_numpy[0,15,:,:] plt.imshow(pre_numpy_person) plt.show()