Skip to content

Instantly share code, notes, and snippets.

@zwade
Created March 22, 2021 18:17
Show Gist options
  • Save zwade/440466e810e72bd99a6bcadef2b10aa9 to your computer and use it in GitHub Desktop.
Save zwade/440466e810e72bd99a6bcadef2b10aa9 to your computer and use it in GitHub Desktop.
Comet ML + dataTap
from comet_ml import Experiment
import torch
import torchvision
import torchvision.transforms as T
from datatap import Api
from datatap.torch import create_dataloader, torch_to_image_annotation
from datatap.utils.print_helpers import pprint
from datatap.metrics import ConfusionMatrix, PrecisionRecallCurve
from datatap.comet import init_experiment, log_validation_proposals
REPOSITORY = "_/aicrowd-food-recognition-challenge"
BATCH_SIZE = 2
COMET_PROJECT = "your project name"
COMET_WORKSPACE = "your workspace name"
def main():
experiment = Experiment(
project_name = COMET_PROJECT,
workspace = COMET_WORKSPACE,
auto_output_logging = False,
)
api = Api()
database = api.get_default_database()
repository = database.get_repository(REPOSITORY)
dataset = repository.get_dataset("latest")
init_experiment(experiment, dataset)
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
classes = list(dataset.template.classes.keys())
classes_with_background = ["__background__"] + classes
class_map = {
cls: i + 1
for i, cls in enumerate(classes)
}
class_map_with_background = {
cls: i
for i, cls in enumerate(classes_with_background)
}
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(
pretrained=False,
pretrained_backbone=True,
num_classes=len(classes_with_background),
).to(device)
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(
params,
lr=0.001,
momentum=0.9,
weight_decay=0.0005
)
lr_scheduler = torch.optim.lr_scheduler.StepLR(
optimizer,
step_size=3,
gamma=0.1
)
image_transform = T.Compose([
T.Resize(100),
T.ToTensor(),
])
training_data = create_dataloader(
dataset,
"training",
batch_size = BATCH_SIZE,
image_transform = image_transform,
device = torch.device("cpu"),
class_mapping = class_map_with_background
)
validation_data = create_dataloader(
dataset,
"validation",
batch_size = BATCH_SIZE,
image_transform = image_transform,
device = "cpu",
class_mapping = class_map_with_background
)
confusion_matrix = ConfusionMatrix(classes = classes)
pr_curve = PrecisionRecallCurve()
num_epochs = 4
for i in range(num_epochs):
model.train()
pprint("Starting epoch: {green}{0}", i)
count = 0
for batch in training_data:
optimizer.zero_grad()
targets = [
{
"boxes": boxes.to(device),
"labels": labels.to(device)
}
for boxes, labels in zip(batch.boxes, batch.labels)
]
images = [image.to(device) for image in batch.images]
losses = model(images, targets)
total_loss = sum([val for val in losses.values()])
total_loss.backward()
optimizer.step()
count += 1
if count % 10 == 0:
pprint(
"Iter: {yellow}{0:4d}{clear}, Classifier loss: {yellow}{1:2.3f}{clear}, Total loss: {yellow}{2:2.3f}{clear}",
count,
losses["loss_classifier"],
total_loss
)
model.eval()
all_annotations = []
for batch in validation_data:
predictions = model([image.to(device) for image in batch.images])
annotations = [
torch_to_image_annotation(
image,
class_map,
labels = prediction["labels"],
scores = prediction["scores"],
boxes = prediction["boxes"],
uid = ground_truth.uid,
)
for image, ground_truth, prediction in zip(batch.images, batch.original_annotations, predictions)
]
all_annotations += annotations
confusion_matrix.batch_add_annotation(
batch.original_annotations,
annotations,
iou_threshold = 0.5,
confidence_threshold = 0.5
)
pr_curve.batch_add_annotation(
batch.original_annotations,
annotations,
iou_threshold = 0.5,
)
log_validation_proposals(experiment, all_annotations)
lr_scheduler.step()
if __name__ == "__main__":
main()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment