def gen_transforms_and_learner(stats: CocoDatasetStats,
train_records: List[BaseRecord],
valid_records: List[BaseRecord],
img_sz=128,
bs=4,
acc_cycs=8,
num_workers=2):
train_tfms = tfms.A.Adapter([
*tfms.A.aug_tfms(
size=img_sz,
presize=img_sz + 128,
shift_scale_rotate=tfms.A.ShiftScaleRotate(
shift_limit=.025, scale_limit=0.025, rotate_limit=9)),
tfms.A.Normalize(mean=stats.chn_means / 255, std=stats.chn_stds / 255)
])
valid_tfms = tfms.A.Adapter(
[*tfms.A.resize_and_pad(img_sz),
tfms.A.Normalize()])
train_ds = Dataset(train_records, train_tfms)
valid_ds = Dataset(valid_records, valid_tfms)
# Using gradient accumulation to process minibatch of 32 images in 8 loops, i.e. 8 images per loop.
# I ran this model w img 512x512x3 on my Dell XPS15 w GTX-1050 with 4GB VRAM, 16GM RAM, ~20min/epoch.
backbone_name = "tf_efficientdet_lite0"
model = efficientdet.model(model_name=backbone_name,
img_size=img_sz,
num_classes=len(stats.lbl2name))
train_dl = efficientdet.train_dl(train_ds,
batch_size=bs,
num_workers=num_workers,
shuffle=True)
valid_dl = efficientdet.valid_dl(valid_ds,
batch_size=bs,
num_workers=max(1, num_workers // 2),
shuffle=False)
monitor_metric = 'COCOMetric'
metrics = [COCOMetric(metric_type=COCOMetricType.bbox)]
save_model_fname = f'{backbone_name}-{img_sz}'
callbacks = [
GradientAccumulation(bs * acc_cycs),
SaveModelDupBestCallback(fname=save_model_fname,
monitor=monitor_metric),
EarlyStoppingCallback(monitor=monitor_metric,
min_delta=0.001,
patience=10),
FastGPUMonitorCallback(delay=1)
]
learn = efficientdet.fastai.learner(dls=[train_dl, valid_dl],
model=model,
metrics=metrics,
cbs=callbacks)
learn.freeze()
return valid_tfms, learn, backbone_name
def fridge_efficientdet_model() -> nn.Module:
WEIGHTS_URL = "https://github.com/airctic/model_zoo/releases/download/m2/fridge_tf_efficientdet_lite0.zip"
model = efficientdet.model("tf_efficientdet_lite0",
num_classes=5,
img_size=512)
state_dict = torch.hub.load_state_dict_from_url(
WEIGHTS_URL, map_location=torch.device("cpu"))
model.load_state_dict(state_dict)
return model