auto-scheduler for pytorch task.
Project description
Fed Flow
Description
auto-scheduler for pytorch task.
Install
pip instal fedflow
Usage
import os
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torchvision.datasets import mnist
from torchvision.transforms import transforms
from fedflow import Task, TaskGroup, FedFlow
from fedflow.utils.trainer.supervised_trainer import SupervisedTrainer
abspath = os.path.abspath(".")
print("Downloading...")
mnist.MNIST(root=os.path.join(abspath, "datasets"),
download=True,
train=True,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.13066062,), (0.30810776,))
]))
print("Download Finished.")
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(1, 20, 5, 1)
self.conv2 = nn.Conv2d(20, 50, 5, 1)
self.fc1 = nn.Linear(4 * 4 * 50, 500)
self.fc2 = nn.Linear(500, 10)
self.softmax = nn.Softmax(dim=1)
def forward(self, x):
x = F.relu(self.conv1(x))
x = F.max_pool2d(x, 2, 2)
x = F.relu(self.conv2(x))
x = F.max_pool2d(x, 2, 2)
x = x.view(-1, 4 * 4 * 50)
x = F.relu(self.fc1(x))
x = self.fc2(x)
return x
class MnistTask(Task):
def __init__(self, id):
super(MnistTask, self).__init__(task_id=id)
def load(self):
self.mnist_dataset = mnist.MNIST(root=os.path.join(abspath, "datasets"),
download=True,
train=True,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.13066062,), (0.30810776,))
]))
self.mnist_model = Net()
self.mnist_optim = optim.Adam(self.mnist_model.parameters(), lr=0.1)
self.criterion = nn.CrossEntropyLoss()
def train(self):
self.mnist_model = self.mnist_model.to(self.device)
trainer = SupervisedTrainer(self.mnist_model, self.mnist_optim, self.criterion, epoch=20, device=self.device,
console_out="console.out")
trainer.mount_dataset(self.mnist_dataset, batch_size=32)
trainer.train()
if __name__ == "__main__":
for i in range(3):
group = TaskGroup()
for j in range(10):
task = MnistTask(10 * i + j)
group.add_task(task)
FedFlow.add_group(group)
FedFlow.start()
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
fedflow-0.2.0b1.tar.gz
(16.0 kB
view hashes)
Built Distribution
fedflow-0.2.0b1-py3-none-any.whl
(20.7 kB
view hashes)
Close
Hashes for fedflow-0.2.0b1-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 60752fde0a69cedef79f18604341b4afd26dc783e6d1b4a869a151f44ea0029c |
|
MD5 | af7fd0018a0fec47540eaf365b7f9047 |
|
BLAKE2b-256 | e8fa122797a20b796d3af28b7f7a464cc918b0e081c8764249484c21b0977e92 |