Refactored ffnn
@@ -223,23 +223,10 @@ class FFNNSynapse(Synapse):
transform = F.relu(F.max_pool2d(self.transform_drop(self.transform_conv2(transform)),2))
transform = transform.view(-1, self.transform_dim)
- # remote_context: responses from a bittensor remote network call.
- # remote_context.shape = [batch_size, bittensor.__network_dim__]
- if remote:
- remote_context, keys, scores = self.call_remote(images, transform)
- output.keys = keys
- output.scores = scores
-
# local_context: distillation model for remote_context.
# local_context.shape = [batch_size, bittensor.__network_dim__]
local_context = self.context_layer1(transform.detach())
local_context = self.context_layer2(local_context)
- if remote:
- # distillation_loss: distillation loss between local_context and remote_context
- # distillation_loss.shape = [1]
- distillation_loss = F.mse_loss(local_context, remote_context.detach())
- output.distillation_loss = distillation_loss
- output.loss = output.loss + distillation_loss
# local_hidden: hidden layer encoding using local_context.
# local_hidden.shape = [batch_size, bittensor.__network_dim__]
@@ -247,6 +234,7 @@ class FFNNSynapse(Synapse):
local_hidden = F.relu(self.hidden_layer1(local_hidden))
local_hidden = F.relu(self.hidden_layer2(local_hidden))
output.local_hidden = local_hidden
+
if targets is not None:
# local_target: projection of local_hidden onto target dimension.
# local_target.shape = [batch_size, target_dim]
@@ -262,15 +250,32 @@ class FFNNSynapse(Synapse):
output.local_target_loss = local_target_loss
output.loss = output.loss + local_target_loss
+ if remote:
+ output = self.forward_remote(local_context, output, images, transform, targets)
+
+ return output
+
+ def forward_remote(self, local_context, output, images, transform, targets):
+ # remote_context: responses from a bittensor remote network call.
+ # remote_context.shape = [batch_size, bittensor.__network_dim__]
+ remote_context, keys, scores = self.call_remote(images, transform)
+ output.keys = keys
+ output.scores = scores
+
+ # distillation_loss: distillation loss between local_context and remote_context
+ # distillation_loss.shape = [1]
+ distillation_loss = F.mse_loss(local_context, remote_context.detach())
+ output.distillation_loss = distillation_loss
+ output.loss = output.loss + distillation_loss
+
# remote_hidden: hidden layer encoding using remote_context.
# remote_hidden.shape = [batch_size, bittensor.__network_dim__]
- if remote:
remote_hidden = torch.cat([transform, remote_context], dim=1)
remote_hidden = self.hidden_layer1(remote_hidden)
remote_hidden = self.hidden_layer2(remote_hidden)
output.remote_hidden = remote_hidden
- if remote and targets is not None:
+ if targets is not None:
# remote_target: projection of remote_hidden onto target dimension.
# remote_target.shape = [batch_size, target_dim]
remote_target = self.target_layer1(remote_hidden)
@@ -285,3 +290,12 @@ class FFNNSynapse(Synapse):
output.remote_target_loss = remote_target_loss
return output
+
+
+
+
+
+
+
+
+
