def __init__(self, options, dataset):
model = choose_model(options)
self.move_model_to_gpu(model, options)
self.optimizer = GD(model.parameters(), lr=options['lr'], weight_decay=options['wd'])
self.num_epoch = options['num_epoch']
worker = LrdWorker(model, self.optimizer, options)
super(FedAvg5Trainer, self).__init__(options, dataset, worker=worker)
def __init__(self, train_data, test_data):
self.model = CifarCnn((3, 32, 32), 10)
self.optimizer = GD(self.model.parameters(), lr=0.1, weight_decay=0.001)
self.batch_size = 64
self.num_epoch = 100
self.train_dataloader = DataLoader(train_data, batch_size=self.batch_size, shuffle=True)
self.test_dataloader = DataLoader(test_data, batch_size=self.batch_size, shuffle=False)
self.criterion = CrossEntropyLoss()
def __init__(self, options, dataset):
model = choose_model(options)
self.move_model_to_gpu(model, options)
self.optimizer = GD(model.parameters(),
lr=options['lr'],
weight_decay=options['wd'])
super(FedAvgTrainer, self).__init__(options, dataset, model,
self.optimizer)
def __init__(self, model, options):
# Basic parameters
self.model = model
self.optimizer = GD(model.parameters(),
lr=options['lr'],
weight_decay=options['wd'])
self.num_epoch = options['num_epoch']
self.lr = options['lr']
self.meta_lr = options['meta_lr']
self.gpu = options['gpu'] if 'gpu' in options else False