联邦元学习算法Per-FedAvg的PyTorch实现

def train(args, model):
    model.train()
    Dtr, Dte, m, n = nn_seq(model.name, args.B)
    model.len = len(Dtr)
    print('training...')
    data = [x for x in iter(Dtr)]
    for epoch in range(args.E):
        model = one_step(args, data, model, lr=args.alpha)
        model = one_step(args, data, model, lr=args.beta)

    return model


def one_step(args, data, model, lr):
    ind = np.random.randint(0, high=len(data), size=None, dtype=int)
    seq, label = data[ind]
    seq = seq.to(args.device)
    label = label.to(args.device)
    y_pred = model(seq)
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    loss_function = nn.MSELoss().to(args.device)
    loss = loss_function(y_pred, label)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    return model

def local_adaptation(args, model):
    model.train()
    Dtr, Dte = nn_seq_wind(model.name, 50)
    optimizer = torch.optim.Adam(model.parameters(), lr=args.alpha)
    loss_function = nn.MSELoss().to(args.device)
    loss = 0
    for epoch in range(1):
        for seq, label in Dtr:
            seq, label = seq.to(args.device), label.to(args.device)
            y_pred = model(seq)
            loss = loss_function(y_pred, label)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        # print('local_adaptation loss', loss.item())

    return model