|
|
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 |
