Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

V0.3 pretrained rep fix #64

Merged
merged 3 commits into from
Jul 3, 2023
Merged

V0.3 pretrained rep fix #64

Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
29ae362
fixed error in base_nets for R3MConv
MBronars Jul 3, 2023
1979437
added error messages for Sequential
MBronars Jul 3, 2023
a4d3352
added doc string for Sequential init
MBronars Jul 3, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
12 changes: 9 additions & 3 deletions robomimic/models/base_nets.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -88,11 +88,15 @@ class Sequential(torch.nn.Sequential, Module):
"""
Compose multiple Modules together (defined above).
"""
def __init__(self, *args):
def __init__(self, *args, has_output_shape = True):
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

add small docstring here please

for arg in args:
assert isinstance(arg, Module)
if has_output_shape:
assert isinstance(arg, Module)
else:
assert isinstance(arg, nn.Module)
torch.nn.Sequential.__init__(self, *args)
self.fixed = False
self.has_output_shape = has_output_shape

def output_shape(self, input_shape=None):
"""
Expand All @@ -106,6 +110,8 @@ def output_shape(self, input_shape=None):
Returns:
out_shape ([int]): list of integers corresponding to output shape
"""
if not self.has_output_shape:
raise NotImplementedError("Output shape is not defined for this module")
out_shape = input_shape
for module in self:
out_shape = module.output_shape(out_shape)
Expand Down Expand Up @@ -574,7 +580,7 @@ def __init__(
transforms.CenterCrop(224),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
)
self.nets = nn.Sequential(*([preprocess] + list(net.module.convnet.children())))
self.nets = Sequential(*([preprocess] + list(net.module.convnet.children())), has_output_shape = False)
if freeze:
self.nets.freeze()

Expand Down