from typing import Dict, Tuple, Union import copy import torch import torch.nn as nn import torchvision import hydra from SUREFlow.utils.crop_randomizer import CropRandomizer class ModuleAttrMixin(nn.Module): def __init__(self): super().__init__() self._dummy_variable = nn.Parameter() @property def device(self): return next(iter(self.parameters())).device @property def dtype(self): return next(iter(self.parameters())).dtype class MultiImageObsEncoder(ModuleAttrMixin): def __init__(self, shape_meta: dict, rgb_model: Union[nn.Module, Dict[str, nn.Module]], resize_shape: Union[Tuple[int, int], Dict[str, tuple], None] = None, crop_shape: Union[Tuple[int, int], Dict[str, tuple], None] = None, random_crop: bool = True, # replace BatchNorm with GroupNorm use_group_norm: bool = False, # use single rgb model for all rgb inputs share_rgb_model: bool = False, # renormalize rgb input with imagenet normalization # assuming input in [0,1] imagenet_norm: bool = False ): """ Assumes rgb input: B,C,H,W Assumes low_dim input: B,D """ super().__init__() rgb_keys = list() key_model_map = nn.ModuleDict() key_transform_map = nn.ModuleDict() key_shape_map = dict() # handle sharing vision backbone if share_rgb_model: assert isinstance(rgb_model, nn.Module) key_model_map['rgb'] = rgb_model obs_shape_meta = shape_meta['obs'] for key, attr in obs_shape_meta.items(): shape = tuple(attr['shape']) type = attr.get('type') key_shape_map[key] = shape if type == 'rgb': rgb_keys.append(key) # configure model for this key this_model = None if not share_rgb_model: this_model = copy.deepcopy(hydra.utils.instantiate(rgb_model)) # if isinstance(rgb_model, DictConfig): # # have provided model for each key # this_model = rgb_model[key] # # if isinstance(rgb_model, DictConfig): # # this_model = hydra.utils.instantiate(rgb_model[key]) # # elif isinstance(rgb_model, dict): # # this_model = rgb_model[key] # else: # assert isinstance(rgb_model, nn.Module) # # have a copy of the rgb model # this_model = copy.deepcopy(rgb_model) if this_model is not None: # if use_group_norm: # this_model = replace_submodules( # root_module=this_model, # predicate=lambda x: isinstance(x, nn.BatchNorm2d), # func=lambda x: nn.GroupNorm( # num_groups=x.num_features // 16, # num_channels=x.num_features) # ) key_model_map[key] = this_model # configure resize input_shape = shape this_resizer = nn.Identity() if resize_shape is not None: if isinstance(resize_shape, dict): h, w = resize_shape[key] else: h, w = resize_shape this_resizer = torchvision.transforms.Resize( size=(h, w) ) input_shape = (shape[0], h, w) # configure randomizer this_randomizer = nn.Identity() if crop_shape is not None: if isinstance(crop_shape, dict): h, w = crop_shape[key] else: h, w = crop_shape if random_crop: this_randomizer = CropRandomizer( input_shape=input_shape, crop_height=h, crop_width=w, num_crops=1, pos_enc=False ) else: this_randomizer = torchvision.transforms.CenterCrop( size=(h, w) ) # configure normalizer this_normalizer = nn.Identity() if imagenet_norm: this_normalizer = torchvision.transforms.Normalize( mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) this_transform = nn.Sequential(this_resizer, this_randomizer, this_normalizer) key_transform_map[key] = this_transform else: raise RuntimeError(f"Unsupported obs type: {type}") rgb_keys = sorted(rgb_keys) self.shape_meta = shape_meta self.key_model_map = key_model_map self.key_transform_map = key_transform_map self.share_rgb_model = share_rgb_model self.rgb_keys = rgb_keys self.key_shape_map = key_shape_map def forward(self, obs_dict, lang_cond=None): batch_size = None features = list() # process rgb input if self.share_rgb_model: # pass all rgb obs to rgb model imgs = list() for key in self.rgb_keys: img = obs_dict[key] if batch_size is None: batch_size = img.shape[0] else: assert batch_size == img.shape[0] assert img.shape[1:] == self.key_shape_map[key] img = self.key_transform_map[key](img) imgs.append(img) # (N*B,C,H,W) imgs = torch.cat(imgs, dim=0) # (N*B,D) feature = self.key_model_map['rgb'](imgs) # (N,B,D) feature = feature.reshape(-1, batch_size, *feature.shape[1:]) # (B,N,D) feature = torch.moveaxis(feature, 0, 1) # (B,N*D) feature = feature.reshape(batch_size, -1) features.append(feature) else: # run each rgb obs to independent models for key in self.rgb_keys: img = obs_dict[key] if batch_size is None: batch_size = img.shape[0] else: assert batch_size == img.shape[0] try: assert img.shape[1:] == self.key_shape_map[key] except AssertionError as e: print(f"key: {key}, shape: {img.shape[1:]}, expected: {self.key_shape_map[key]}") raise e #change image from B H W C -> B C H W # img = img.permute(0, 3, 1, 2) has been changed in dataset img = self.key_transform_map[key](img) if lang_cond is None: feature = self.key_model_map[key](img) else: feature = self.key_model_map[key](img, lang_cond) features.append(feature) # concatenate all features result = torch.stack(features, dim=1) # result = torch.cat(features, dim=-1) return result # shape-> [B, f1 + f2] -> [B, 64 + 8] @torch.no_grad() def output_shape(self): example_obs_dict = dict() obs_shape_meta = self.shape_meta['obs'] batch_size = 1 for key, attr in obs_shape_meta.items(): shape = tuple(attr['shape']) this_obs = torch.zeros( (batch_size,) + shape, dtype=self.dtype, device=self.device) example_obs_dict[key] = this_obs example_output = self.forward(example_obs_dict) output_shape = example_output.shape[1:] return output_shape