def train(self, model: nn.Module, train_data: DataContainer,
context: FederatedLearning.Context,
config: TrainerParams) -> Tuple[any, int]:
model.to(self.device)
model.train()
optimizer = config.get_optimizer()(model)
criterion = config.get_criterion()
epoch_loss = []
for epoch in range(config.epochs):
batch_loss = []
for batch_idx, (x, labels) in enumerate(
train_data.batch(config.batch_size)):
x = x.to(self.device)
labels = labels.to(self.device)
optimizer.zero_grad()
log_probs = model(x)
loss = criterion(log_probs, labels)
loss.backward()
optimizer.step()
batch_loss.append(loss.item())
if len(batch_loss) > 0:
epoch_loss.append(sum(batch_loss) / len(batch_loss))
weights = model.cpu().state_dict()
return weights, len(train_data)
def infer(self, model: nn.Module, test_data: DataContainer):
model.to(self.device)
model.eval()
test_loss = test_acc = test_total = 0.
criterion = self.criterion
with torch.no_grad():
for batch_idx, (x, target) in enumerate(
test_data.batch(self.batch_size)):
x = x.to(self.device)
target = target.to(self.device)
pred = model(x)
loss = criterion(pred, target)
_, predicted = torch.max(pred, -1)
correct = predicted.eq(target).sum()
test_acc += correct.item()
test_loss += loss.item() * target.size(0)
test_total += target.size(0)
return test_acc / test_total, test_loss / test_total