dodox/dppo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

support varying img size

Browse Source
This commit is contained in:
allenzren
2024-09-16 17:55:31 -04:00
parent 64595baca9
commit 1aaa6c2302
18 changed files with 131 additions and 81 deletions
Download Patch File Download Diff File Expand all files Collapse all files
model/common/critic.py
+4 -5
View File
@@ -122,7 +122,6 @@ class ViTCritic(CriticObs):
cond_dim,
img_cond_steps=1,
spatial_emb=128,
patch_repr_dim=128,
dropout=0,
augment=False,
num_img=1,
@@ -136,8 +135,8 @@ class ViTCritic(CriticObs):
self.img_cond_steps = img_cond_steps
if num_img > 1:
self.compress1 = SpatialEmb(
num_patch=121, # TODO: repr_dim // patch_repr_dim,
patch_dim=patch_repr_dim,
num_patch=self.backbone.num_patch,
patch_dim=self.backbone.patch_repr_dim,
prop_dim=cond_dim,
proj_dim=spatial_emb,
dropout=dropout,
@@ -145,8 +144,8 @@ class ViTCritic(CriticObs):
self.compress2 = deepcopy(self.compress1)
else: # TODO: clean up
self.compress = SpatialEmb(
num_patch=121,
patch_dim=patch_repr_dim,
num_patch=self.backbone.num_patch,
patch_dim=self.backbone.patch_repr_dim,
prop_dim=cond_dim,
proj_dim=spatial_emb,
dropout=dropout,
model/common/mlp_gaussian.py
+5 -7
View File
@@ -32,8 +32,6 @@ class Gaussian_VisionMLP(nn.Module):
std_max=1,
spatial_emb=0,
visual_feature_dim=128,
repr_dim=96 * 96,
patch_repr_dim=128,
dropout=0,
num_img=1,
augment=False,
@@ -51,8 +49,8 @@ class Gaussian_VisionMLP(nn.Module):
assert spatial_emb > 1, "this is the dimension"
if num_img > 1:
self.compress1 = SpatialEmb(
num_patch=121, # TODO: repr_dim // patch_repr_dim,
patch_dim=patch_repr_dim,
num_patch=self.backbone.num_patch,
patch_dim=self.backbone.patch_repr_dim,
prop_dim=cond_dim,
proj_dim=spatial_emb,
dropout=dropout,
@@ -60,8 +58,8 @@ class Gaussian_VisionMLP(nn.Module):
self.compress2 = deepcopy(self.compress1)
else: # TODO: clean up
self.compress = SpatialEmb(
num_patch=121,
patch_dim=patch_repr_dim,
num_patch=self.backbone.num_patch,
patch_dim=self.backbone.patch_repr_dim,
prop_dim=cond_dim,
proj_dim=spatial_emb,
dropout=dropout,
@@ -69,7 +67,7 @@ class Gaussian_VisionMLP(nn.Module):
visual_feature_dim = spatial_emb * num_img
else:
self.compress = nn.Sequential(
nn.Linear(repr_dim, visual_feature_dim),
nn.Linear(self.backbone.repr_dim, visual_feature_dim),
nn.LayerNorm(visual_feature_dim),
nn.Dropout(dropout),
nn.ReLU(),
model/common/vit.py
+37 -10
View File
@@ -3,12 +3,13 @@ ViT image encoder implementation from IBRL, https://github.com/hengyuan-hu/ibrl
"""
from dataclasses import dataclass, field
from dataclasses import dataclass
from typing import List
import einops
import torch
from torch import nn
from torch.nn.init import trunc_normal_
import math
@dataclass
@@ -29,6 +30,8 @@ class VitEncoder(nn.Module):
obs_shape: List[int],
cfg: VitEncoderConfig,
num_channel=3,
img_h=96,
img_w=96,
):
super().__init__()
self.obs_shape = obs_shape
@@ -40,8 +43,11 @@ class VitEncoder(nn.Module):
num_head=cfg.num_heads,
depth=cfg.depth,
num_channel=num_channel,
img_h=img_h,
img_w=img_w,
)
self.img_h = img_h
self.img_w = img_w
self.num_patch = self.vit.num_patches
self.patch_repr_dim = self.cfg.embed_dim
self.repr_dim = self.cfg.embed_dim * self.vit.num_patches
@@ -56,11 +62,11 @@ class VitEncoder(nn.Module):
class PatchEmbed1(nn.Module):
def __init__(self, embed_dim, num_channel=3):
def __init__(self, embed_dim, num_channel=3, img_h=96, img_w=96):
super().__init__()
self.conv = nn.Conv2d(num_channel, embed_dim, kernel_size=8, stride=8)
self.num_patch = 144
self.num_patch = math.ceil(img_h / 8) * math.ceil(img_w / 8)
self.patch_dim = embed_dim
def forward(self, x: torch.Tensor):
@@ -70,7 +76,7 @@ class PatchEmbed1(nn.Module):
class PatchEmbed2(nn.Module):
def __init__(self, embed_dim, use_norm, num_channel=3):
def __init__(self, embed_dim, use_norm, num_channel=3, img_h=96, img_w=96):
super().__init__()
layers = [
nn.Conv2d(num_channel, embed_dim, kernel_size=8, stride=4),
@@ -80,7 +86,11 @@ class PatchEmbed2(nn.Module):
]
self.embed = nn.Sequential(*layers)
self.num_patch = 121 # TODO: specifically for 96x96 set by Hengyuan?
H1 = math.ceil((img_h - 8) / 4) + 1
W1 = math.ceil((img_w - 8) / 4) + 1
H2 = math.ceil((H1 - 3) / 2) + 1
W2 = math.ceil((W1 - 3) / 2) + 1
self.num_patch = H2 * W2
self.patch_dim = embed_dim
def forward(self, x: torch.Tensor):
@@ -146,14 +156,25 @@ class MinVit(nn.Module):
num_head,
depth,
num_channel=3,
img_h=96,
img_w=96,
):
super().__init__()
if embed_style == "embed1":
self.patch_embed = PatchEmbed1(embed_dim, num_channel=num_channel)
self.patch_embed = PatchEmbed1(
embed_dim,
num_channel=num_channel,
img_h=img_h,
img_w=img_w,
)
elif embed_style == "embed2":
self.patch_embed = PatchEmbed2(
embed_dim, use_norm=embed_norm, num_channel=num_channel
embed_dim,
use_norm=embed_norm,
num_channel=num_channel,
img_h=img_h,
img_w=img_w,
)
else:
assert False
@@ -233,8 +254,14 @@ def test_transformer_layer():
if __name__ == "__main__":
obs_shape = [6, 96, 96]
enc = VitEncoder([6, 96, 96], VitEncoderConfig())
obs_shape = [6, 128, 128]
enc = VitEncoder(
obs_shape,
VitEncoderConfig(),
num_channel=obs_shape[0],
img_h=obs_shape[1],
img_w=obs_shape[2],
)
print(enc)
x = torch.rand(1, *obs_shape) * 255
model/diffusion/mlp_diffusion.py
+5 -7
View File
@@ -34,8 +34,6 @@ class VisionDiffusionMLP(nn.Module):
residual_style=False,
spatial_emb=0,
visual_feature_dim=128,
repr_dim=96 * 96,
patch_repr_dim=128,
dropout=0,
num_img=1,
augment=False,
@@ -53,8 +51,8 @@ class VisionDiffusionMLP(nn.Module):
assert spatial_emb > 1, "this is the dimension"
if num_img > 1:
self.compress1 = SpatialEmb(
num_patch=121, # TODO: repr_dim // patch_repr_dim,
patch_dim=patch_repr_dim,
num_patch=self.backbone.num_patch,
patch_dim=self.backbone.patch_repr_dim,
prop_dim=cond_dim,
proj_dim=spatial_emb,
dropout=dropout,
@@ -62,8 +60,8 @@ class VisionDiffusionMLP(nn.Module):
self.compress2 = deepcopy(self.compress1)
else: # TODO: clean up
self.compress = SpatialEmb(
num_patch=121,
patch_dim=patch_repr_dim,
num_patch=self.backbone.num_patch,
patch_dim=self.backbone.patch_repr_dim,
prop_dim=cond_dim,
proj_dim=spatial_emb,
dropout=dropout,
@@ -71,7 +69,7 @@ class VisionDiffusionMLP(nn.Module):
visual_feature_dim = spatial_emb * num_img
else:
self.compress = nn.Sequential(
nn.Linear(repr_dim, visual_feature_dim),
nn.Linear(self.backbone.repr_dim, visual_feature_dim),
nn.LayerNorm(visual_feature_dim),
nn.Dropout(dropout),
nn.ReLU(),