summaryrefslogtreecommitdiff
path: root/models/rgb_part_net.py
diff options
context:
space:
mode:
authorJordan Gong <jordan.gong@protonmail.com>2021-04-04 17:44:23 +0800
committerJordan Gong <jordan.gong@protonmail.com>2021-04-04 17:44:23 +0800
commit6f3dd9109b8ae7b37e3373d844a6c406d83c2b35 (patch)
treea530221dfef3100a236c4091c3d0c15ea636d9e5 /models/rgb_part_net.py
parent6a8824e4fb8bdd1f3e763b78b765830788415cfb (diff)
parent85627d4cfb495453a7c28b3f131b84b1038af674 (diff)
Merge branch 'disentangling_only' into disentangling_only_py3.8disentangling_only_py3.8
Diffstat (limited to 'models/rgb_part_net.py')
-rw-r--r--models/rgb_part_net.py15
1 files changed, 10 insertions, 5 deletions
diff --git a/models/rgb_part_net.py b/models/rgb_part_net.py
index 1c7a1a2..6be6b0a 100644
--- a/models/rgb_part_net.py
+++ b/models/rgb_part_net.py
@@ -10,6 +10,7 @@ from models.auto_encoder import AutoEncoder
class RGBPartNet(nn.Module):
def __init__(
self,
+ num_class: int,
ae_in_channels: int = 3,
ae_in_size: Tuple[int, int] = (64, 48),
ae_feature_channels: int = 64,
@@ -22,11 +23,15 @@ class RGBPartNet(nn.Module):
self.image_log_on = image_log_on
self.ae = AutoEncoder(
- ae_in_channels, ae_in_size, ae_feature_channels, f_a_c_p_dims
+ num_class,
+ ae_in_channels,
+ ae_in_size,
+ ae_feature_channels,
+ f_a_c_p_dims
)
- def forward(self, x_c1, x_c2=None):
- losses, features, images = self._disentangle(x_c1, x_c2)
+ def forward(self, x_c1, x_c2=None, y=None):
+ losses, features, images = self._disentangle(x_c1, x_c2, y)
if self.training:
losses = torch.stack(losses)
@@ -34,11 +39,11 @@ class RGBPartNet(nn.Module):
else:
return features
- def _disentangle(self, x_c1_t2, x_c2_t2=None):
+ def _disentangle(self, x_c1_t2, x_c2_t2=None, y=None):
n, t, c, h, w = x_c1_t2.size()
if self.training:
x_c1_t1 = x_c1_t2[:, torch.randperm(t), :, :, :]
- ((f_a_, f_c_, f_p_), losses) = self.ae(x_c1_t2, x_c1_t1, x_c2_t2)
+ ((f_a_, f_c_, f_p_), losses) = self.ae(x_c1_t2, x_c1_t1, x_c2_t2, y)
f_a = f_a_.view(n, t, -1)
f_c = f_c_.view(n, t, -1)
f_p = f_p_.view(n, t, -1)