class ResNet50_FasterRCNN: def __init__(self, pretrained=False): # Building our FasterRCNN model for objects detection backbone = resnet_fpn_backbone('resnet50', pretrained=pretrained) num_classes = 4 + 1 anchor_generator = AnchorGenerator(sizes=(40, 60, 150, 200, 250), aspect_ratios=(0.7, 1.0, 1.3)) self.model = FRCNN(backbone, num_classes=num_classes, rpn_anchor_generator=anchor_generator) def train(self): self.model.train() def to(self, device): self.model.to(device) def eval(self): self.model.eval() def parameters(self): return self.model.parameters() def get_state_dict(self): return self.model.state_dict() def set_state_dict(self, state_dict): self.model.load_state_dict(state_dict) def fit_batch(self, images, target): return self.model(images, target) def predict_batch(self, images): return self.model(images)
def train_one_epoch(train_data_loader: DataLoader, model: FasterRCNN, device: str, logger: Logger, optimizer: SGD, args: argparse.Namespace) -> Tuple[float, float]: """ Training the model for one epoch :param train_data_loader: Dataloader for training data :param model: FasterRCNN model :param device: Device for the computation :param logger: Logger for logging handling :param optimizer: Optmizer for the training :param args: Arguments :return: """ start = time.time() model.train() loss_per_iteration = [] for idx, result in enumerate(train_data_loader): images = list(image.to(device) for image in result[0]) targets = result[1] targets = [{k: v.to(device) for k, v in t.items()} for t in targets] optimizer.zero_grad() with torch.set_grad_enabled(True): loss_dict = model(images, targets) losses = sum(loss for loss in loss_dict.values()) losses.backward() optimizer.step() loss_per_iteration.append(losses) loss_doc_str = "".join("{}:{} ".format(key, val) for key, val in loss_dict.items()) if idx % args.print_status == 0: logger.info(f'Iteration: [{idx}/{len(train_data_loader)}]\t' f'Loss: {losses} \t' f'Loss_dict: {loss_doc_str}') epoch_time = time.time() - start return loss_per_iteration, epoch_time
def train(args): # use_cuda = args.num_gpus > 0 # logger.debug("Number of gpus available - {}".format(args.num_gpus)) # kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {} device = torch.device( "cuda") if torch.cuda.is_available() else torch.device("cpu") # set the seed for generating random numbers torch.manual_seed(args.seed) if device.type == 'cuda': torch.cuda.manual_seed(args.seed) train_loader = _get_train_data_loader(args.resize) test_loader = _get_test_data_loader(args.resize) logger.debug("Processes {}/{} ({:.0f}%) of train data".format( len(train_loader.sampler), len(train_loader.dataset), 100. * len(train_loader.sampler) / len(train_loader.dataset))) model = FasterRCNN().to(device) params = [p for p in model.parameters() if p.requires_grad] optimizer = optim.Adam(params, lr=args.lr) for epoch in range(1, args.epochs + 1): print("Epoch: ", epoch) model.train() for batch_idx, (batch_images, batch_targets) in enumerate(train_loader): images = list(img.to(device) for img in batch_images) targets = [{k: v.to(device) for k, v in t.items()} for t in batch_targets] loss_dict = model(images, targets) losses = sum(loss for loss in loss_dict.values()) model.zero_grad() losses.backward() optimizer.step() if batch_idx % args.log_interval == 0: logger.info( 'Train Epoch: {} [{}/{} ({:.0f}%)] Loss: {:.6f}'.format( epoch, batch_idx * len(images), len(train_loader.sampler), 100. * batch_idx / len(train_loader), losses.item())) save_model(model, args.model_directory, args.fn)
def evaluate_for_losses(model: FasterRCNN, data_loader: DataLoader, device: torch.device, max_n_batches: int = 1) -> Dict[str, float]: model = model.to(device) model.train() loss_dicts = [] with torch.no_grad(): for i, (img, labels) in enumerate(tqdm(data_loader)): if i >= max_n_batches: break img = list(image.to(device) for image in img) labels = [{k: v.to(device) for k, v in t.items()} for t in labels] loss_dicts.append(model(img, labels)) return { k: np.mean([ld[k].cpu().item() for ld in loss_dicts]) for k in loss_dicts[0].keys() }
def get_fasterrcnn_model(arch_str, num_classes, pretrained=True, pretrained_backbone=True, trainable_layers=5, **kwargs): """Creates FasterRCNN model with resnet backbone""" #if pretrained == True: pretrained_backbone=False backbone = resnet_fpn_backbone(arch_str, pretrained=pretrained_backbone, trainable_layers=trainable_layers) anchor_sizes = ( (16, ), (32, ), (64, ), (128, ), (256, ), ) aspect_ratios = ((0.5, 1.0, 2.0), ) * len(anchor_sizes) anchor_generator = AnchorGenerator(sizes=anchor_sizes, aspect_ratios=aspect_ratios) model = FasterRCNN( backbone, num_classes=num_classes, rpn_anchor_generator=anchor_generator, box_fg_iou_thresh=0.5, box_bg_iou_thresh=0.5, image_mean=[0.0, 0.0, 0.0], # already normalized by fastai image_std=[1.0, 1.0, 1.0], #min_size = 1, #box_score_thresh=0.6, **kwargs) if pretrained: try: pretrained_dict = load_state_dict_from_url( _model_urls['fasterrcnn_' + arch_str + '_fpn_coco'], progress=True) model_dict = model.state_dict() pretrained_dict = { k: v for k, v in pretrained_dict.items() if (k in model_dict) and ( model_dict[k].shape == pretrained_dict[k].shape) } model_dict.update(pretrained_dict) model.load_state_dict(model_dict) #overwrite_eps(model, 0.0) for module in model.modules(): if isinstance(module, FrozenBatchNorm2d): module.eps = 0.0 except Exception as e: #print(e) print("No pretrained coco model found for fasterrcnn_" + arch_str) print("This does not affect the backbone.") return model.train()
def get_fasterrcnn_model_swin(arch_str, num_classes, pretrained=False, pretrained_backbone=True, **kwargs): """Creates FasterRCNN model with swin transformer backbone""" anchor_sizes = ( (32, ), (64, ), (128, ), (256, ), ) aspect_ratios = ((0.5, 1.0, 2.0), ) * len(anchor_sizes) anchor_generator = AnchorGenerator(sizes=anchor_sizes, aspect_ratios=aspect_ratios) #roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0','1','2','3'], # output_size=7, # sampling_ratio=2) img_size = 224 if arch_str in "swin_tiny swin_small".split() else 384 window_size = 7 if arch_str in "swin_tiny swin_small".split() else 12 depths = [2, 2, 6, 2] if arch_str == "swin_tiny" else [2, 2, 18, 2] scale_factors = { "swin_tiny": 1.0, "swin_small": 1.5, "swin_base": 2.0, "swin_large": 2.0 } sf = scale_factors[arch_str] embed_dim = int(96 * sf) fpn_cin = [int(96 * sf * 2**i) for i in range(4)] #fpn_cin = [int(i*sf) for i in [96, 192, 384, 768]] backbone = SwinTransformerFPN(img_size=img_size, window_size=window_size, embed_dim=embed_dim, depths=depths, fpn_cin=fpn_cin, fpn_cout=256) if pretrained_backbone: sd = load_state_dict_from_url(_model_urls[f'{arch_str}_{img_size}'], progress=True, map_location=default_device())['model'] sd_model = backbone.state_dict() sd = {k: v for k, v in sd.items() if k in sd_model.keys()} sd_model.update(sd) backbone.load_state_dict(sd_model) model = FasterRCNN( backbone, num_classes=num_classes, rpn_anchor_generator=anchor_generator, #box_roi_pool=roi_pooler, box_fg_iou_thresh=0.5, box_bg_iou_thresh=0.5, image_mean=[0.0, 0.0, 0.0], # already normalized by fastai image_std=[1.0, 1.0, 1.0], #min_size=IMG_SIZE, #max_size=IMG_SIZE, **kwargs) return model.train()
class FasterRCNNMODEL: #TODO: Later on enable passing params params def __init__(self, model_params=None): self.params = model_params self.model = None self.optimizer = None self.device = torch.device( 'cuda') if torch.cuda.is_available() else torch.device('cpu') def set_backbone(self, backbone): """ backbone is a string containing the backbone we want to use in the model. add more options """ if 'vgg' in backbone.lower(): "to somthing-check for options" elif 'mobilenet_v2' in backbone.lower(): self.backbone = torchvision.models.mobilenet_v2( pretrained=True).features self.backbone.out_channels = 1280 elif 'resnet50' in backbone.lower(): self.backbone = torchvision.models.resnet50( pretrained=True).features self.backbone.out_channels = 256 def set_model(self): """ Set model and determine configuration :return: None, generate self.model to be used for training and testing """ # Default values: box_score_thresh = 0.05, box_nms_thresh = 0.5 kwargs = { 'box_score_thresh': 0.3, 'box_nms_thresh': 0.3, 'box_detections_per_img': 6 } # self.model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=False, # pretrained_backbone=True, # **kwargs) self.model = FasterRCNN(self.backbone, num_classes=7, **kwargs) device = torch.device( 'cuda') if torch.cuda.is_available() else torch.device('cpu') num_classes = 7 in_features = self.model.roi_heads.box_predictor.cls_score.in_features self.model.roi_heads.box_predictor = FastRCNNPredictor( in_features, num_classes) # Allow Multiple GPUs: # if torch.cuda.device_count() > 1: # self.model = nn.DataParallel(self.model) self.model = self.model.to(device) if self.params is None: params = [p for p in self.model.parameters() if p.requires_grad] else: # TODO: Enable user defined model params pass self.optimizer = torch.optim.SGD(params, lr=0.01) def train_model(self, train_loader, num_epochs): """ Train (only!) of the model :param train_loader: DataLoader object :param num_epochs: int. Number of epochs to train the model :return: None, """ self.model.train() # Set to training mode for epoch in range(num_epochs): for images, targets in train_loader: images = list(image.to(self.device) for image in images) targets = [{k: v.to(self.device) for k, v in t.items()} for t in targets] # Zero Gradients self.optimizer.zero_grad() # self.model = self.model.double() # Calculate Loss loss_dict = self.model(images, targets) # what happens here? losses = sum(loss for loss in loss_dict.values()) losses.backward() # Update weights self.optimizer.step() print('Train Loss = {:.4f}'.format(losses.item())) def train_eval_model(self, train_loader, val_loader, num_epochs): """ Train model and evaluate performance after each epoch :param train_loader: DataLoader object. Training images and targets :param val_loader: DataLoader object. validation images and targets :param num_epochs: int. Number of epochs for training and validation :return: """ # For evaluation imgs_name_list = [] bbox_list = [] labels_list = [] for epoch in range(num_epochs): train_loss = 0 val_loss = 0 self.model.train() # Set to training mode with torch.set_grad_enabled(True): for images, targets in train_loader: # Pass data to GPU images = list(image.to(self.device) for image in images) targets = [{k: v.to(self.device) for k, v in t.items()} for t in targets] # Zero Gradients self.optimizer.zero_grad() # self.model = self.model.double() # Calculate Loss loss_dict = self.model(images, targets) # what happens here? losses = sum(loss for loss in loss_dict.values()) train_loss += losses.item() * len(images) # Backward Prop & Update weights losses.backward() self.optimizer.step() print('Train Loss = {:.4f}'.format(train_loss / len(train_loader.dataset))) # TODO: Calculate Dice and IoU loss for it with torch.no_grad(): for idx, (imgs_name, images, targets) in enumerate(val_loader): self.model.train() images = list(image.to(self.device) for image in images) targets = [{k: v.to(self.device) for k, v in t.items()} for t in targets] loss_dict = self.model(images, targets) losses = sum(loss for loss in loss_dict.values()) val_loss += losses.item() * len(images) if epoch == num_epochs - 1: self.model.eval() # Set model to evaluate performance targets = self.model(images) # Think of moving all this into gen_out_file - Looks nicer imgs_name_list.extend(imgs_name) bbox_list.extend([ target['boxes'].int().cpu().tolist() for target in targets ]) labels_list.extend([ target['labels'].int().cpu().tolist() for target in targets ]) """Optional - SEE the performance on the second last batch""" if (epoch == num_epochs - 1) and idx == (len(val_loader) - 2): self.model.eval() # Set model to evaluate performance targets = self.model(images) MiscUtils.view(images, targets, k=len(images), model_type='faster_rcnn') DataUtils.gen_out_file('output_file.txt', imgs_name_list, bbox_list, labels_list) print('Validation Loss = {:.4f}'.format( val_loss / len(val_loader.dataset)))
model.to(DEVICE) optimizer = optim.Adam(model.parameters(), lr=1e-5) train_dataset = Dataset() test_dataset = Dataset(training=False) train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, pin_memory=True, drop_last=True) test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False, drop_last=False) for epoch in range(EPOCH): for index, (data, target) in enumerate(train_loader): model.train() data = data.to(DEVICE) target['boxes'] = target['boxes'].to(DEVICE) target['labels'] = target['labels'].squeeze(0).to(DEVICE) optimizer.zero_grad() dic = [] boxes = target['boxes'] labels = target['labels'] for i in range(len(labels)): dic.append({'boxes': boxes[i].unsqueeze(0), 'labels': labels[i]}) output = model(data, dic) loss = sum(l for l in output.values()) if index % 5 == 0: print("Epoch {}[{:.1f}%]\tLoss: {:.6f}".format( epoch + 1, 100 * index / len(train_loader), loss.item())) loss.backward()