class FFNNSynapse(bittensor.synapse.Synapse): """ Simple feed forward NN for images. """ def __init__(self, config: Munch, **kwargs): r""" Init a new ffnn synapse module. :param [config]: munch namespace config item. :type [config]: [:obj:`munch.Munch`](, `required`) """ super(FFNNSynapse, self).__init__(config=config, **kwargs) if config == None: config = FFNNSynapse.default_config() bittensor.config.Config.update_with_kwargs(config.synapse, kwargs) FFNNSynapse.check_config(config) self.config = config # transform_layer: transforms images to common dimension. # [batch_size, -1, -1, -1] -> [batch_size, self.transform_dim] self.transform = Normalize((0.1307, ), (0.3081, ), device=self.device) self.transform_pool = nn.AdaptiveAvgPool2d((28, 28)) self.transform_conv1 = nn.Conv2d(1, 10, kernel_size=5) self.transform_conv2 = nn.Conv2d(10, 20, kernel_size=5) self.transform_drop = nn.Dropout2d() self.transform_dim = 320 # context_layer: distills the remote_context from the transform layer. # [batch_size, transform_dim] -> [batch_size, bittensor.__network_dim__] self.context_layer1 = nn.Linear(self.transform_dim, 256) self.context_layer2 = nn.Linear(256, bittensor.__network_dim__) # hidden_layer: learns hidden units for network and target. # [batch_size, transform_dim + bittensor.__network_dim__] = [batch_size, bittensor.__network_dim__] self.hidden_layer1 = nn.Linear( self.transform_dim + bittensor.__network_dim__, bittensor.__network_dim__) self.hidden_layer2 = nn.Linear(bittensor.__network_dim__, bittensor.__network_dim__) # dendrite: (PKM layer) queries network using pooled embeddings as context. # [batch_size, -1] -> topk * [batch_size, bittensor.__network_dim__] self.router = PKMRouter(config, query_dim=bittensor.__network_dim__) # target_layer: Maps from hidden layer to target dimension # [batch_size, bittensor.__network_dim__] -> [batch_size, self.target_dim] self.target_layer1 = nn.Linear(bittensor.__network_dim__, 256) self.target_layer2 = nn.Linear(256, self.config.synapse.target_dim) self.to(self.device) @staticmethod def default_config() -> Munch: parser = argparse.ArgumentParser() FFNNSynapse.add_args(parser) config = bittensor.config.Config.to_config(parser) return config @staticmethod def add_args(parser: argparse.ArgumentParser): parser.add_argument( '--synapse.target_dim', default=10, type=int, help='Final logit layer dimension. i.e. 10 for MNIST.') parser = PKMRouter.add_args(parser) @staticmethod def check_config(config: Munch): assert config.synapse.target_dim > 0, "target dimension must be greater than 0." config = PKMRouter.check_config(config) def forward_image(self, images: torch.Tensor): r""" Forward image inputs through the FFNN synapse . Args: inputs (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_dim, channels, rows, cols)`, `required`): Image tensors produced by calling PIL.toTensor() and with sequence dimension. Returns: hidden (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_dim, bittensor.__network_dim__)`, `required`): Hidden layer encoding produced by using local_context. """ # images: remove sequence dimension from images. # images.shape = [batch_size, channels, rows, cols] images = images.view(images.shape[0] * images.shape[1], images.shape[2], images.shape[3], images.shape[4]).to(self.device) # hidden: hidden layer using local_contextcontext for local computation only. # hidden.shape = [batch_size, __network_dim__] hidden = self.local_forward(images=images).local_hidden # hidden: re-add sequence dimension to outputs. # hidden.shape = [batch_size, sequence_dim, __network_dim__] hidden = torch.unsqueeze(hidden, 1) return hidden def local_forward(self, images: torch.Tensor, targets: torch.Tensor = None) -> SimpleNamespace: r""" Forward pass non-sequential image inputs and targets through the FFNN Synapse. The call does not make remote queries to the network and returns only local hidden, target and losses. Args: images (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, channels, rows, cols)`, `required`): PIL.toTensor() encoded images. targets (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, target_dim)`, `optional`, defaults to None): Image labels. Returns: local_context (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, bittensor.__network_dim__)`, `required`): Pre-Hidden layer context, trained to match the remote context. local_hidden (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, bittensor.__network_dim__)`, `required`): Hidden layer produced from the context. local_target (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, target_dim)`, `optional`): FFNN Target predictions using local_context. local_target_loss (:obj:`torch.FloatTensor` of shape :obj:`(1)`, `optional`): FFNN Classification loss using local_context. local_accuracy (:obj:`torch.FloatTensor` of shape :obj:`(1)`, `optional`): Accuracy of target predictions. """ # Return vars to be filled. output = SimpleNamespace() # transform: transform images to common shape. # transform.shape = [batch_size, self.transform_dim] transform = self.transform(images).to(self.device) transform = F.relu(F.max_pool2d(self.transform_conv1(transform), 2)) transform = F.relu( F.max_pool2d(self.transform_drop(self.transform_conv2(transform)), 2)) output.transform = transform.view(-1, self.transform_dim) # local_context: distillation model for remote_context. # local_context.shape = [batch_size, bittensor.__network_dim__] local_context = self.context_layer1(output.transform.detach()) output.local_context = self.context_layer2(local_context) # local_hidden: hidden layer encoding using local_context. # local_hidden.shape = [batch_size, bittensor.__network_dim__] local_hidden = torch.cat( (output.transform, output.local_context.detach()), dim=1) local_hidden = F.relu(self.hidden_layer1(local_hidden)) output.local_hidden = F.relu(self.hidden_layer2(local_hidden)) if targets is not None: # local_target: projection of local_hidden onto target dimension. # local_target.shape = [batch_size, target_dim] targets.to(self.device) local_target = self.target_layer1(output.local_hidden) local_target = self.target_layer2(local_target) output.local_target = F.log_softmax(local_target, dim=1) # local_target_loss: loss between local_target and passed targets. # local_target_loss.shape = [1] output.local_target_loss = F.nll_loss(output.local_target, targets) # Record extra metadata accuracy. max_logit = local_target.data.max(1, keepdim=True)[1] correct = max_logit.eq(targets.data.view_as(max_logit)).sum() output.local_accuracy = (100.0 * correct) / targets.shape[0] return output def remote_forward(self, neuron: bittensor.neuron.Neuron, images: torch.Tensor, targets: torch.Tensor = None) -> SimpleNamespace: """ Forward pass non-sequential image inputs and targets through the remote context of the synapse. The call makes RPC queries accross the network using the passed neuron's metagraph and dendrite. Args: neuron (:obj: `bittensor.neuron.Neuron`, `required`): Bittensor neuron, used for making queries to the remote network. images (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, channels, rows, cols)`, `required`): PIL.toTensor() encoded images. targets (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, target_dim)`, `optional`, defaults to None): Image labels. Returns: self.local_forward() + SimpleNamespace ( router (:obj:`SimpleNamespace`, `required`): Outputs from the pkm dendrite remote call. distillation_loss (:obj:`torch.FloatTensor` of shape :obj:`(1)`, `optional`): Distillation loss between the local and remote context. remote_hidden (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, bittensor.__network_dim__)`, `optional`): Hidden layer encoding produced using the remote context. remote_target (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, target_dim)`, `optional`): FFNN Target predictions using the remote_context. remote_target_loss (:obj:`torch.FloatTensor` of shape :obj:`(1)`, `optional`): FFNN Classification loss using the remote_context. """ # Call the local forward pass. # output = bittensor.SynapseOutput output = self.local_forward(images, targets) # Make remote queries using the PKMRouter. # remote_context: responses from a bittensor remote network call. # remote_context.shape = [batch_size, bittensor.__network_dim__] images = torch.unsqueeze(images, 1) output.router = self.router.forward_image(neuron, images, output.local_hidden) remote_context = torch.squeeze(output.router.response, 1).to(self.device) # Distill the local context to match the remote context. # distillation_loss: distillation loss between local_context and remote_context # distillation_loss.shape = [1] output.distillation_loss = F.mse_loss(output.local_context, remote_context.detach()) # remote_hidden: hidden layer encoding using remote_context. # remote_hidden.shape = [batch_size, bittensor.__network_dim__] remote_hidden = torch.cat([output.transform, remote_context], dim=1) remote_hidden = self.hidden_layer1(remote_hidden) output.remote_hidden = self.hidden_layer2(remote_hidden) if targets is not None: # Project hidden units onto the targets. # remote_target: projection of remote_hidden onto target dimension. # remote_target.shape = [batch_size, target_dim] remote_target = self.target_layer1(remote_hidden) remote_target = self.target_layer2(remote_target) output.remote_target = F.log_softmax(remote_target, dim=1) # Compute the target loss. # remote_target_loss: loss between remote_target and passed targets. # remote_target_loss.shape = [1] output.remote_target_loss = F.nll_loss(output.remote_target, targets) # Add extra metrics # Record extra metadata accuracy. max_logit = output.remote_target.data.max(1, keepdim=True)[1] correct = max_logit.eq(targets.data.view_as(max_logit)).sum() output.remote_accuracy = (100.0 * correct) / targets.shape[0] return output
class DPNSynapse(bittensor.synapse.Synapse): """ Bittensor endpoint trained on PIL images to detect objects using an DPN. """ def __init__( self, config: Munch = None): r""" Init a new DPN synapse module. Args: config (:obj: `munch.Munch`, `required`) munch namespace config item. """ super(DPNSynapse, self).__init__(config = config) if config == None: config = DPNSynapse.build_config() in_planes, out_planes = config.synapse.in_planes, config.synapse.out_planes num_blocks, dense_depth = config.synapse.num_blocks, config.synapse.dense_depth # Transform Network """ Transform network. Layers take in image inputs normalizes them and applies 4 convolutional layers. Image encoder: transforms PIL-encoded tensors to a common shape. [batch_size, channels, rows, cols] -> [batch_size, -1, -1, -1] Output: [batch_size, self.transform_dim (9728)] """ self.transform = Normalize((0.1307,), (0.3081,), device=self.device) self.adaptive_pool = nn.AdaptiveAvgPool2d((32, 32)) self.transform_conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False) self.transform_bn1 = nn.BatchNorm2d(64) self.last_planes = 64 self.transform_layer1 = self._make_layer(in_planes[0], out_planes[0], num_blocks[0], dense_depth[0], stride=1) self.transform_layer2 = self._make_layer(in_planes[1], out_planes[1], num_blocks[1], dense_depth[1], stride=2) self.transform_layer3 = self._make_layer(in_planes[2], out_planes[2], num_blocks[2], dense_depth[2], stride=1) self.transform_layer4 = self._make_layer(in_planes[3], out_planes[3], num_blocks[3], dense_depth[3], stride=2) self.transform_dim = (out_planes[3] * 4)+(((num_blocks[3]+1) * 4)*dense_depth[3]) # dendrite: (PKM layer) queries network using pooled embeddings as context. # [batch_size, -1] -> topk * [batch_size, bittensor.__network_dim__] self.router = PKMRouter(config, query_dim = self.transform_dim) # Context layers. """ Distillation model for remote context. This layer takes input coming from transform layer, and runs it through 3 linear layers, projecting it to bittensor.__network_dim__. """ self.context_layer1 = nn.Linear(self.transform_dim, 512) self.context_layer2 = nn.Linear(512, 256) self.context_layer3 = nn.Linear(256, bittensor.__network_dim__) # hidden layer. self.hidden_layer1 = nn.Linear(self.transform_dim + bittensor.__network_dim__, 512) self.hidden_layer2 = nn.Linear(512, 256) self.hidden_layer3 = nn.Linear(256, bittensor.__network_dim__) # Layers to project target down to target size passed by config # (number of classes) self.target_layer1 = nn.Linear(bittensor.__network_dim__, 128) self.target_layer2 = nn.Linear(128, self.config.synapse.target_dim) self.to(self.device) @staticmethod def build_config() -> Munch: parser = argparse.ArgumentParser(); DPNSynapse.add_args(parser) config = bittensor.config.Config.to_config(parser); DPNSynapse.check_config(config) return config @staticmethod def add_args(parser: argparse.ArgumentParser): r""" This function adds the configuration items for the DPN synapse. These args are use to instantiate a Dual Path model. Instantiating a configuration with the defaults will yield a "shallow" DPN-26 configuration. For deeper network configurations, it is possible to set the num_blocks parameter to (3, 4, 20, 3) for a DPN-92. For DPN-98 set the following: in_planes: (160, 320, 640, 1280) out_planes: (256, 512, 1024, 2048) num_blocks: (3, 6, 20, 3) dense_depth: (16, 32, 32, 128) """ def to_list(arg): return [int(i) for i in arg.split(",")] parser.add_argument('--synapse.in_planes', default='160, 320, 640, 1280', action="append", type=to_list) parser.add_argument('--synapse.out_planes', default='256, 512, 1024, 2048', action="append", type=to_list) parser.add_argument('--synapse.num_blocks', default='3, 6, 20, 3', action="append", type=to_list) parser.add_argument('--synapse.dense_depth', default='16, 32, 32, 128', action="append", type=to_list) parser.add_argument('--synapse.target_dim', default=10, type=int, help='Final logit layer dimension. i.e. 10 for CIFAR-10.') parser = PKMRouter.add_args(parser) @staticmethod def check_config(config: Munch): assert isinstance(config.synapse.in_planes, list), 'synapse.in_planes must be a tuple, got {}'.format(config.synapse.in_planes) assert isinstance(config.synapse.out_planes, list), 'synapse.out_planes must be a tuple, got {}'.format(config.synapse.out_planes) assert isinstance(config.synapse.num_blocks, list), 'synapse.num_blocks must be a tuple, got {}'.format(config.synapse.num_blocks) assert isinstance(config.synapse.dense_depth, list), 'synapse.dense_depth must be a tuple, got {}'.format(config.synapse.dense_depth) assert all(isinstance(el, int) for el in config.synapse.in_planes), 'synapse.in_planes must be a tuple of ints, got {}'.format(config.synapse.in_planes) assert all(isinstance(el, int) for el in config.synapse.out_planes), 'synapse.out_planes must be a tuple of ints, got {}'.format(config.synapse.out_planes) assert all(isinstance(el, int) for el in config.synapse.num_blocks), 'synapse.num_blocks must be a tuple of ints, got {}'.format(config.synapse.num_blocks) assert all(isinstance(el, int) for el in config.synapse.dense_depth), 'synapse.dense_depth must be a tuple of ints, got {}'.format(config.synapse.dense_depth) def forward_image ( self, images: torch.Tensor): r""" Forward image inputs through the DPN synapse . Args: inputs (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_dim, channels, rows, cols)`, `required`): Image tensors produced by calling PIL.toTensor() and with sequence dimension. Returns: hidden (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_dim, bittensor.__network_dim__)`, `required`): Hidden layer encoding produced by using local_context. """ # images: remove sequence dimension from images. # images.shape = [batch_size, channels, rows, cols] images = images.view(images.shape[0] * images.shape[1], images.shape[2], images.shape[3], images.shape[4]) # hidden: hidden layer using local context for local computation only. # hidden.shape = [batch_size, __network_dim__] hidden = self.forward (images = images.to(self.device), remote = False).local_hidden # hidden: re-add sequence dimension to outputs. # hidden.shape = [batch_size, sequence_dim, __network_dim__] hidden = torch.unsqueeze(hidden, 1) return hidden def local_forward ( self, images: torch.Tensor, targets: torch.Tensor = None ) -> SimpleNamespace: r""" Forward pass non-sequential image inputs and targets through the DPN Synapse. Args: images (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, channels, rows, cols)`, `required`): PIL.toTensor() encoded images. targets (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, config.target_size)`, `optional`): Image labels. remote (:obj:`bool')`, `optional`): Switch between local and remote context. If true, function makes quries to the remote network. Returns: SimpleNamespace ( local_context (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, bittensor.__network_dim__)`, `required`): Pre-Hidden layer context, trained to match the remote context. local_hidden (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, bittensor.__network_dim__)`, `required`): Hidden layer produced from the context. local_target (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, target_dim)`, `optional`): FFNN Target predictions using local_context. local_target_loss (:obj:`torch.FloatTensor` of shape :obj:`(1)`, `optional`): FFNN Classification loss using local_context. local_accuracy (:obj:`torch.FloatTensor` of shape :obj:`(1)`, `optional`): Accuracy of target predictions. transform (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, transform_dim)`, `optional`): transformation of various sized images to batch-size transform dim. ) """ # Return vars to be filled. output = SimpleNamespace () r""" Transform the images into a common shape (32x32) """ # transform: transform images to common shape. # transform.shape = [batch_size, self.transform_dim] transform = self.transform(images) transform = self.adaptive_pool(transform) transform = F.relu(self.transform_bn1(self.transform_conv1(transform.detach()))) transform = self.transform_layer1(transform) transform = self.transform_layer2(transform) transform = self.transform_layer3(transform) transform = self.transform_layer4(transform) transform = F.avg_pool2d(transform, 4) output.transform = torch.flatten(transform, start_dim=1) # local_context: distillation model for remote_context. # local_context.shape = [batch_size, bittensor.__network_dim__] local_context = self.context_layer1(output.transform.detach()) local_context = self.context_layer2(local_context) output.local_context = self.context_layer3(local_context) # local_hidden: hidden layer encoding using local_context. # local_hidden.shape = [batch_size, bittensor.__network_dim__] local_hidden = torch.cat([output.transform, output.local_context], dim=1) local_hidden = self.hidden_layer1(local_hidden) local_hidden = self.hidden_layer2(local_hidden) output.local_hidden = self.hidden_layer3(local_hidden) if targets is not None: # local_target: projection of local_hidden onto target dimension. # local_target.shape = [batch_size, target_dim] targets.to(self.device) local_target = self.target_layer1(output.local_hidden) local_target = self.target_layer2(local_target) output.local_target = F.log_softmax(local_target, dim=1) # local_target_loss: loss between local_target and passed targets. # local_target_loss.shape = [1] output.local_target_loss = F.nll_loss(output.local_target, targets) # Record extra metadata accuracy. max_logit = local_target.data.max(1, keepdim=True)[1] correct = max_logit.eq( targets.data.view_as(max_logit) ).sum() output.local_accuracy = (100.0 * correct) / targets.shape[0] return output def remote_forward(self, neuron: bittensor.neuron.Neuron, images: torch.Tensor, targets: torch.Tensor = None) -> SimpleNamespace: """ Forward pass non-sequential image inputs and targets through the synapse. Makes RPC queries to downstream neurons. Args: neuron (:obj: `bittensor.neuron.Neuron`, `required`): Bittensor neuron, used for making queries to the remote network. images (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, channels, rows, cols)`, `required`): PIL.toTensor() encoded images. targets (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, target_dim)`, `optional`, defaults to None): Image labels. Returns: self.local_forward() + SimpleNamespace ( router (:obj:`SimpleNamespace`, `required`): Outputs from the pkm dendrite remote call. distillation_loss (:obj:`torch.FloatTensor` of shape :obj:`(1)`, `optional`): Distillation loss between the local and remote context. remote_hidden (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, bittensor.__network_dim__)`, `optional`): Hidden layer encoding produced using the remote context. remote_target (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, target_dim)`, `optional`): FFNN Target predictions using the remote_context. remote_target_loss (:obj:`torch.FloatTensor` of shape :obj:`(1)`, `optional`): FFNN Classification loss using the remote_context. ) """ # Call the local forward pass. # output = bittensor.SynapseOutput output = self.local_forward( images, targets ) # Make remote queries using the PKMRouter. # remote_context: responses from a bittensor remote network call. # remote_context.shape = [batch_size, bittensor.__network_dim__] images = torch.unsqueeze(images, 1) output.router = self.router.forward_image( neuron, images, output.transform ) remote_context = torch.squeeze( output.router.response, 1 ).to(self.device) # Distill the local context to match the remote context. # distillation_loss: distillation loss between local_context and remote_context # distillation_loss.shape = [1] output.distillation_loss = F.mse_loss(output.local_context, remote_context.detach() ) # remote_hidden: hidden layer encoding using remote_context. # remote_hidden.shape = [batch_size, bittensor.__network_dim__] remote_hidden = torch.cat([output.transform, remote_context], dim=1) remote_hidden = self.hidden_layer1(remote_hidden) remote_hidden = self.hidden_layer2(remote_hidden) output.remote_hidden = self.hidden_layer3(remote_hidden) if targets is not None: # remote_target: projection of remote_hidden onto target dimension. # remote_target.shape = [batch_size, config.target_size] remote_target = self.target_layer1(output.remote_hidden) remote_target = self.target_layer2(remote_target) output.remote_target = F.log_softmax(remote_target, dim=1) # remote_target_loss: loss between remote_target and passed targets. # remote_target_loss.shape = [1] output.remote_target_loss = F.nll_loss(output.remote_target, targets) return output def _make_layer(self, in_planes, out_planes, num_blocks, dense_depth, stride): """ Generates a sequential container containing Bottleneck layers. Args: in_planes (tuple): 4-element tuple describing the in_planes config. out_planes (tuple): 4-element tuple describing the out_planes config. num_blocks (tuple): 4-element tuple describing the number of blocks at this layer. dense_depth (tuple): 4-element tuple describing the depth of this layer. stride (int): Convolutional stride length. Returns: nn.Sequential: A torch.nn sequential container containing the layers outlined in the inputs. """ strides = [stride] + [1]*(num_blocks-1) layers = [] for i,stride in enumerate(strides): layers.append(self.Bottleneck(self.last_planes, in_planes, out_planes, dense_depth, stride, i==0)) self.last_planes = out_planes + (i+2) * dense_depth return nn.Sequential(*layers) class Bottleneck(nn.Module): def __init__(self, last_planes, in_planes, out_planes, dense_depth, stride, first_layer): super(DPNSynapse.Bottleneck, self).__init__() self.out_planes = out_planes self.dense_depth = dense_depth self.conv1 = nn.Conv2d(last_planes, in_planes, kernel_size=1, bias=False) self.bn1 = nn.BatchNorm2d(in_planes) self.conv2 = nn.Conv2d(in_planes, in_planes, kernel_size=3, stride=stride, padding=1, groups=32, bias=False) self.bn2 = nn.BatchNorm2d(in_planes) self.conv3 = nn.Conv2d(in_planes, out_planes+dense_depth, kernel_size=1, bias=False) self.bn3 = nn.BatchNorm2d(out_planes + dense_depth) self.shortcut = nn.Sequential() if first_layer: self.shortcut = nn.Sequential( nn.Conv2d(last_planes, out_planes + dense_depth, kernel_size=1, stride=stride, bias=False), nn.BatchNorm2d(out_planes + dense_depth) ) def forward(self, x): out = F.relu(self.bn1(self.conv1(x))) out = F.relu(self.bn2(self.conv2(out))) out = self.bn3(self.conv3(out)) x = self.shortcut(x) d = self.out_planes out = torch.cat([x[:,:d,:,:]+out[:,:d,:,:], x[:,d:,:,:], out[:,d:,:,:]], 1) out = F.relu(out) return out