| | import torch
|
| | import torch.nn as nn
|
| | from typing import Dict, Any, List
|
| | import asyncio
|
| | import websockets
|
| | import json
|
| | from pydantic import BaseModel
|
| |
|
| | class PeerMessage(BaseModel):
|
| | message_type: str
|
| | payload: Dict[str, Any]
|
| | peer_id: str
|
| |
|
| | class DecentModel(nn.Module):
|
| | """Base class for decentralized deep learning models"""
|
| |
|
| | def __init__(self):
|
| | super().__init__()
|
| | self.peer_id = self._generate_peer_id()
|
| | self.peers: List[str] = []
|
| | self.websocket = None
|
| | self.state_updates = {}
|
| |
|
| | def _generate_peer_id(self) -> str:
|
| | """Generate a unique peer ID"""
|
| | import uuid
|
| | return str(uuid.uuid4())
|
| |
|
| | async def connect_to_network(self, network_url: str):
|
| | """Connect to the decentralized network"""
|
| | self.websocket = await websockets.connect(network_url)
|
| | await self._register_peer()
|
| |
|
| | async def _register_peer(self):
|
| | """Register this peer with the network"""
|
| | message = PeerMessage(
|
| | message_type="register",
|
| | payload={"model_type": self.__class__.__name__},
|
| | peer_id=self.peer_id
|
| | )
|
| | await self.websocket.send(message.json())
|
| |
|
| | async def broadcast_state_update(self, state_dict: Dict[str, torch.Tensor]):
|
| | """Broadcast model state updates to other peers"""
|
| | message = PeerMessage(
|
| | message_type="state_update",
|
| | payload={"state": self._serialize_state_dict(state_dict)},
|
| | peer_id=self.peer_id
|
| | )
|
| | await self.websocket.send(message.json())
|
| |
|
| | def _serialize_state_dict(self, state_dict: Dict[str, torch.Tensor]) -> Dict[str, List[float]]:
|
| | """Serialize model state for transmission"""
|
| | return {k: v.cpu().numpy().tolist() for k, v in state_dict.items()}
|
| |
|
| | async def receive_state_updates(self):
|
| | """Receive and process state updates from other peers"""
|
| | while True:
|
| | message = await self.websocket.recv()
|
| | data = PeerMessage.parse_raw(message)
|
| | if data.message_type == "state_update":
|
| | self.state_updates[data.peer_id] = self._deserialize_state_dict(
|
| | data.payload["state"]
|
| | )
|
| |
|
| | def _deserialize_state_dict(self, state_dict: Dict[str, List[float]]) -> Dict[str, torch.Tensor]:
|
| | """Deserialize received model state"""
|
| | return {k: torch.tensor(v) for k, v in state_dict.items()}
|
| |
|
| | def aggregate_states(self):
|
| | """Aggregate state updates from all peers"""
|
| | if not self.state_updates:
|
| | return
|
| |
|
| |
|
| | aggregated_state = {}
|
| | for key in self.state_updates[list(self.state_updates.keys())[0]].keys():
|
| | tensors = [states[key] for states in self.state_updates.values()]
|
| | aggregated_state[key] = torch.mean(torch.stack(tensors), dim=0)
|
| |
|
| |
|
| | self.load_state_dict(aggregated_state)
|
| | self.state_updates.clear()
|
| |
|
| | def forward(self, *args, **kwargs):
|
| | """Forward pass - to be implemented by child classes"""
|
| | raise NotImplementedError |
