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()