TorchVisionODAdapter does not convert BoxList to array/tensor when using torchvision and pytorch lightning
#1768
Answered
by
AdeelH
aerotractjack
asked this question in
Q&A
-
|
Hello, I'm back again. Trying to use rastervision object detection datasets with a torchvision object detection model and pytorch lightning training. My guess is that I am not applying the import albumentations as A
from rastervision.core.data import ClassConfig
from rastervision.pytorch_learner import (
ObjectDetectionRandomWindowGeoDataset,
ObjectDetectionSlidingWindowGeoDataset,
)
from torch.utils.data import DataLoader
from torchvision.models.detection import fasterrcnn_resnet50_fpn_v2
import pytorch_lightning as pl
import torch
import torch.functional as F
from rastervision.pipeline.file_system import make_dir
from pytorch_lightning.loggers import TensorBoardLogger
from rastervision.pytorch_learner.object_detection_utils import TorchVisionODAdapter
class ObjectDetection(pl.LightningModule):
def __init__(self, backbone, lr=1e-4):
super().__init__()
self.backbone = backbone
self.lr = lr
def forward(self, img):
return self.backbone(img)['out']
def training_step(self, batch, batch_idx):
img, mask = batch
img = img.float()
mask = mask.long()
out = self.forward(img)
loss = F.cross_entropy(out, mask)
log_dict = {'train_loss': loss}
self.log_dict(log_dict, on_step=True, on_epoch=True, prog_bar=True, logger=True)
return loss
def validation_step(self, batch, batch_idx):
img, mask = batch
img = img.float()
mask = mask.long()
out = self.forward(img)
loss = F.cross_entropy(out, mask)
log_dict = {'validation_loss': loss}
self.log_dict(log_dict, on_step=True, on_epoch=True, prog_bar=True, logger=True)
return loss
def configure_optimizers(self):
optimizer = torch.optim.Adam(
self.parameters(), lr=self.lr)
return optimizerAnd here is the code for my training loop class RVLightning:
def __init__(self, tr_uris, val_uris, pred_uris, output, class_config, kw=None):
self.train_uris = tr_uris
self.val_uris = val_uris
self.pred_uris = pred_uris
self.cc = ClassConfig(
names=class_config["names"],
colors=class_config["colors"],
null_class="null")
self.output_uri = output.get("uri")
self.bucket = output.get("bucket")
self.kw = kw
def build_train_ds(self):
kw = self.kw.get("train_ds_kw", {})
train_ds = ObjectDetectionRandomWindowGeoDataset.from_uris(
class_config=self.cc,
image_uri=self.train_uris[0],
aoi_uri=self.train_uris[2],
label_vector_uri=self.train_uris[1],
label_vector_default_class_id=self.cc.get_class_id('DF'),
size_lims=kw.get("size_lims", [300,350]),
out_size=kw.get("out_size", 325),
max_windows=kw.get("max_windows", 10),
)
return train_ds
def build_val_ds(self):
kw = self.kw.get("val_ds_kw", {})
val_ds = ObjectDetectionSlidingWindowGeoDataset.from_uris(
class_config=self.cc,
image_uri=self.val_uris[0],
aoi_uri=self.val_uris[2],
label_vector_uri=self.val_uris[1],
label_vector_default_class_id=self.cc.get_class_id('DF'),
size=kw.get("size", 325),
stride=kw.get("stride", 325))
return val_ds
def build_train_val_loader(self):
tds, vds = self.build_train_ds(), self.build_val_ds()
batch_size = 8
train = DataLoader(tds, batch_size=batch_size, shuffle=True, num_workers=4)
val = DataLoader(vds, batch_size=batch_size, num_workers=4)
return train, val
def train(self):
kw = self.kw.get("train_kw", {})
batch_size = kw.get("batch_size", 8)
lr = float(kw.get("lr", 1e-4))
epochs = kw.get("epochs", 1)
output_dir = self.output_uri
make_dir(output_dir)
fast_dev_run = False
# option 1
backbone = fasterrcnn_resnet50_fpn_v2(num_classes=len(self.cc))
# option 2
model = ObjectDetection(backbone, lr=lr)
tb_logger = TensorBoardLogger(save_dir=output_dir + "/tensorboard", flush_secs=10)
# option 3
trainer = pl.Trainer(
accelerator='auto',
min_epochs=1,
max_epochs=epochs+1,
default_root_dir=output_dir + "/trainer",
logger=[tb_logger],
fast_dev_run=fast_dev_run,
log_every_n_steps=1,
)
train_dl, val_dl = self.build_train_val_loader()
trainer.fit(model, train_dl, val_dl)
trainer.save_checkpoint(output_dir + "/trainer/final-model.ckpt")
def run(config_path):
from configreader import yaml2dict
conf = yaml2dict(config_path)
obj = RVLightning(
conf["train_uri"],
conf["val_uri"],
conf["pred_uri"],
conf["output"],
conf["class_config"],
conf
)
obj.train()Here is the error message I get after trying to train It seems to me I have 3 options for where to apply the ODAdapter. I can't seem to find an example of rastervision+lightning for object detection, and the only code snippets I've found for the ODAdapter were being used with a different pipeline configuration. Is it possible to use the objectdetection*datasets with lightning? |
Beta Was this translation helpful? Give feedback.
Answered by
AdeelH
Apr 5, 2023
Replies: 1 comment 5 replies
-
|
Option 1 is correct. The model doesn't seem to be the problem here though. It's just that you need to pass in a different collate function to the from rastervision.pytorch_learner.object_detection utils import collate_fn |
Beta Was this translation helpful? Give feedback.
5 replies
-
|
You might also want to be careful about the num_classes. Torchvision models expect class ID = 0 to be a background class. In RV OD examples we don't specify a background class in the class config, but RV automatically adds it when creating the model and |
Beta Was this translation helpful? Give feedback.
-
Thank you, I noticed something along those lines last week so since then I have been defining my class with "null" as ID=0 |
Beta Was this translation helpful? Give feedback.
-
|
So, I added in option 1, but cannot get a working training loop. I keep running into issues that are slightly confusing to me. Normally using lightning this is how I would construct a training loop, and the Here is my updated training loop: def train(self):
kw = self.kw.get("train_kw", {})
lr = float(kw.get("lr", 1e-4))
epochs = kw.get("epochs", 1)
output_dir = self.output_uri
make_dir(output_dir)
fast_dev_run = False
backbone = fasterrcnn_resnet50_fpn_v2(
num_classes=len(self.cc), pretrained=True)
# create adapter here
backbone = TorchVisionODAdapter(backbone)
model = ObjectDetection(backbone, lr=lr)
tb_logger = TensorBoardLogger(save_dir=output_dir + "/tensorboard", flush_secs=10)
trainer = pl.Trainer(
accelerator='auto',
min_epochs=1,
max_epochs=epochs+1,
default_root_dir=output_dir + "/trainer",
logger=[tb_logger],
fast_dev_run=fast_dev_run,
log_every_n_steps=1,
)
train_dl, val_dl = self.build_train_val_loader()
trainer.fit(model, train_dl, val_dl)
trainer.save_checkpoint(output_dir + "/trainer/final-model.ckpt")And here is my current model implementation class ObjectDetection(pl.LightningModule):
def __init__(self, backbone, lr=1e-4):
super().__init__()
self.backbone = backbone
self.lr = lr
def forward(self, img):
return self.backbone(img)
def training_step(self, batch, batch_idx):
image, target = batch
loss_dict = self.backbone(image, target)
losses = sum(loss for loss in loss_dict.values())
batch_size = len(batch[0])
self.log_dict(loss_dict, batch_size=batch_size)
self.log("train_loss", losses, batch_size=batch_size)
return losses
def validation_step(self, batch, batch_idx):
image, target = batch
loss_dict = self.backbone(image, target)
losses = sum(loss for loss in loss_dict.values())
batch_size = len(batch[0])
self.log_dict(loss_dict, batch_size=batch_size)
self.log("val_loss", losses, batch_size=batch_size)
return losses
def configure_optimizers(self):
optimizer = torch.optim.Adam(
self.parameters(), lr=self.lr)
return optimizerHere is the error I receive: This is all happening within a sanity check done by torchvision by calling my |
Beta Was this translation helpful? Give feedback.
-
|
Even with a small example like the following, I am getting issues because my ground truth input and model output are predicting different numbers of boxes for an image, and the loss function is complaining about it. My sanity check: def training_step(self, batch, batch_idx):
print("sanity training")
image, target = batch
pred = self.backbone(image, target)
print(pred)Error output: |
Beta Was this translation helpful? Give feedback.
-
|
Sorry, you answered my original question. I'm going to mark this as answered and move my other question to a new thread. Appreciate the help a ton. |
Beta Was this translation helpful? Give feedback.
Answer selected by
aerotractjack
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Option 1 is correct. The model doesn't seem to be the problem here though. It's just that you need to pass in a different collate function to the
DataLoaders instead of using the default one. This one: