def getGradient(batch): inputs = batch['input'] inputs = torch.unsqueeze(inputs, 1) targets = batch['target'] #print(inputs.shape, targets.shape) if args.mixup: inputs, targets = mixup(inputs, targets, num_classes=len(CLASSES)) inputs = Variable(inputs, requires_grad=True) targets = Variable(targets, requires_grad=False) if use_gpu: inputs = inputs.cuda() targets = targets.cuda(async=True) outputs = model(inputs) if args.mixup: loss = mixup_cross_entropy_loss(outputs, targets) else: loss = criterion(outputs, targets) optimizer.zero_grad() loss.backward() this_loss = loss.item() pred = outputs.data.max(1, keepdim=True)[1] if args.mixup: targets = batch['target'] targets = Variable(targets, requires_grad=False).cuda(async=True) this_correct = pred.eq(targets.data.view_as(pred)).sum() this_total = targets.size(0) current_client_grad = torch.zeros(1,1).cuda() shape_list = [] for name, param in model.named_parameters(): if param.requires_grad: #print(name, param.grad.shape, param.grad.type())#, param.grad) #current_client_grad.append(param.grad) shape_list.append(list(param.grad.shape)) current_client_grad = torch.cat((current_client_grad, param.grad.view(-1,1)), 0) #break current_client_grad = current_client_grad[1:,:].view(-1,) return current_client_grad, this_loss, this_correct, this_total, shape_list
def main(): # 1. load dataset, train and valid train_dataset, valid_dataset = build_dataset( n_mels=n_mels, train_dataset=args.train_dataset, valid_dataset=args.valid_dataset, background_noise=args.background_noise) print('train ', len(train_dataset), 'val ', len(valid_dataset)) weights = train_dataset.make_weights_for_balanced_classes() sampler = WeightedRandomSampler(weights, len(weights)) train_dataloader = DataLoader(train_dataset, batch_size=args.batch_size, sampler=sampler, pin_memory=use_gpu, num_workers=args.dataload_workers_nums) valid_dataloader = DataLoader(valid_dataset, batch_size=args.batch_size, shuffle=False, pin_memory=use_gpu, num_workers=args.dataload_workers_nums) # a name used to save checkpoints etc. # 2. prepare the model, checkpoint full_name = '%s_%s_%s_bs%d_lr%.1e_wd%.1e' % ( args.model, args.optim, args.lr_scheduler, args.batch_size, args.learning_rate, args.weight_decay) if args.comment: full_name = '%s_%s' % (full_name, args.comment) model = models.create_model(model_name=args.model, num_classes=len(CLASSES), in_channels=1) if use_gpu: model = torch.nn.DataParallel(model).cuda() criterion = torch.nn.CrossEntropyLoss() if args.optim == 'sgd': optimizer = torch.optim.SGD(model.parameters(), lr=args.learning_rate, momentum=0.9, weight_decay=args.weight_decay) else: optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay) start_timestamp = int(time.time() * 1000) start_epoch = 0 best_accuracy = 0 best_loss = 1e100 global_step = 0 if args.resume: print("resuming getShapeLista checkpoint '%s'" % args.resume) checkpoint = torch.load(args.resume) model.load_state_dict(checkpoint['state_dict']) model.float() optimizer.load_state_dict(checkpoint['optimizer']) best_accuracy = checkpoint.get('accuracy', best_accuracy) best_loss = checkpoint.get('loss', best_loss) start_epoch = checkpoint.get('epoch', start_epoch) global_step = checkpoint.get('step', global_step) del checkpoint # reduce memory if args.lr_scheduler == 'plateau': lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer, patience=args.lr_scheduler_patience, factor=args.lr_scheduler_gamma) else: lr_scheduler = torch.optim.lr_scheduler.StepLR( optimizer, step_size=args.lr_scheduler_step_size, gamma=args.lr_scheduler_gamma, last_epoch=start_epoch - 1) def get_lr(): return optimizer.param_groups[0]['lr'] writer = SummaryWriter(comment=('_speech_commands_' + full_name)) #3. train and validation print("training %s for Google speech commands..." % args.model) since = time.time() #grad_client_list = [[]] * args.clients federated = Federated(args.clients, args.matrix_size, args.num_threads, args.gradient_res_path) for epoch in range(start_epoch, args.max_epochs): print("epoch %3d with lr=%.02e" % (epoch, get_lr())) phase = 'train' writer.add_scalar('%s/learning_rate' % phase, get_lr(), epoch) model.train() # Set model to training mode running_loss = 0.0 it = 0 correct = 0 total = 0 #compute for each client current_client = 0 pbar = tqdm(train_dataloader, unit="audios", unit_scale=train_dataloader.batch_size, disable=False) # for i in range(0, train_dataloader.__len__, self.batch_size): # pbar. #print(pbar.total, len(train_dataloader), len(train_dataset), train_dataloader.batch_size) #num_batches = len(pbar.total) for batch in pbar: inputs = batch['input'] inputs = torch.unsqueeze(inputs, 1) targets = batch['target'] #print(inputs.shape, targets.shape) if args.mixup: inputs, targets = mixup(inputs, targets, num_classes=len(CLASSES)) inputs = Variable(inputs, requires_grad=True) targets = Variable(targets, requires_grad=False) if use_gpu: inputs = inputs.cuda() targets = targets.cuda(async=True) outputs = model(inputs) if args.mixup: loss = mixup_cross_entropy_loss(outputs, targets) else: loss = criterion(outputs, targets) optimizer.zero_grad() loss.backward() current_client_grad = torch.zeros(1, 1).cuda() shape_list = [] for name, param in model.named_parameters(): if param.requires_grad: #print(name, param.grad.shape, param.grad.type())#, param.grad) #current_client_grad.append(param.grad) shape_list.append(list(param.grad.shape)) current_client_grad = torch.cat( (current_client_grad, param.grad.view(-1, 1)), 0) #break current_client_grad = current_client_grad[1:, :].view(-1, ) #print(current_client_grad.shape) if (current_client == 0): federated.init(current_client_grad, shape_list) #print("client ", current_client, " start") start_time = time.time() federated.work_for_client(current_client, current_client_grad) #print("client", current_client, " complete") end_time = time.time() #print("work for client ", current_client, " cost ", end_time - start_time) if (current_client == args.clients - 1): recover_start = time.time() recovered_grad = federated.recoverGradient() ind = 0 recover_end = time.time() #print("recover gradient cost ", recover_end - recover_start) #print(recovered_grad_in_cuda, recovered_grad_in_cuda[0].shape, r) for name, param in model.named_parameters(): if param.requires_grad: param.grad = recovered_grad[ind] ind += 1 assert (ind == len(recovered_grad)) optimizer.step() #print("all clients finished") current_client = 0 else: current_client += 1 # only update the parameters when current_client == args.clients - 1 # statistics it += 1 global_step += 1 #running_loss += loss.data[0] running_loss += loss.item() pred = outputs.data.max(1, keepdim=True)[1] if args.mixup: targets = batch['target'] targets = Variable(targets, requires_grad=False).cuda(async=True) correct += pred.eq(targets.data.view_as(pred)).sum() total += targets.size(0) writer.add_scalar('%s/loss' % phase, loss.item(), global_step) # update the progress bar pbar.set_postfix({ 'loss': "%.05f" % (running_loss / it), 'acc': "%.02f%%" % (100 * float(correct) / total) }) #print("[batch]\t", it, " [loss]\t ", running_loss / it, " [acc] \t", 100 * float(correct)/total) #print('------------------------------------------------------------------') #break accuracy = float(correct) / total epoch_loss = running_loss / it writer.add_scalar('%s/accuracy' % phase, 100 * accuracy, epoch) writer.add_scalar('%s/epoch_loss' % phase, epoch_loss, epoch) if args.lr_scheduler == 'plateau': lr_scheduler.step(epoch_loss) if (accuracy > best_accuracy): best_accuracy = accuracy checkpoint = { 'epoch': epoch, 'step': global_step, 'state_dict': model.state_dict(), 'loss': epoch_loss, 'accuracy': accuracy, 'optimizer': optimizer.state_dict(), } torch.save( checkpoint, 'checkpoints/federated-best-loss-speech-commands-checkpoint-%s.pth' % full_name) torch.save( model, '%d-%s-federated-best-loss.pth' % (start_timestamp, full_name)) del checkpoint federated.dump()
def train(epoch): global global_step print("epoch %3d with lr=%.02e" % (epoch, get_lr())) phase = 'train' writer.add_scalar('%s/learning_rate' % phase, get_lr(), epoch) model.train() # Set model to training mode running_loss = 0.0 it = 0 correct = 0 total = 0 pbar = tqdm(train_dataloader, unit="audios", unit_scale=train_dataloader.batch_size) for batch in pbar: inputs = batch['input'] inputs = torch.unsqueeze(inputs, 1) targets = batch['target'] if args.mixup: inputs, targets = mixup(inputs, targets, num_classes=len(CLASSES)) inputs = Variable(inputs, requires_grad=True) targets = Variable(targets, requires_grad=False) if use_gpu: inputs = inputs.cuda() targets = targets.cuda(async=True) # forward/backward outputs = model(inputs) if args.mixup: loss = mixup_cross_entropy_loss(outputs, targets) else: loss = criterion(outputs, targets) optimizer.zero_grad() loss.backward() optimizer.step() # statistics it += 1 global_step += 1 running_loss += loss.item() pred = outputs.data.max(1, keepdim=True)[1] if args.mixup: targets = batch['target'] targets = Variable(targets, requires_grad=False).cuda(async=True) correct += pred.eq(targets.data.view_as(pred)).sum() total += targets.size(0) writer.add_scalar('%s/loss' % phase, loss.item(), global_step) # update the progress bar pbar.set_postfix({ 'loss': "%.05f" % (running_loss / it), 'acc': "%.02f%%" % (100 * correct / total) }) accuracy = correct / total epoch_loss = running_loss / it writer.add_scalar('%s/accuracy' % phase, 100 * accuracy, epoch) writer.add_scalar('%s/epoch_loss' % phase, epoch_loss, epoch)
def main(): # 1. load dataset, train and valid train_dataset, valid_dataset = build_dataset( n_mels=n_mels, train_dataset=args.train_dataset, valid_dataset=args.valid_dataset, background_noise=args.background_noise) print('train ', len(train_dataset), 'val ', len(valid_dataset)) weights = train_dataset.make_weights_for_balanced_classes() sampler = WeightedRandomSampler(weights, len(weights)) train_dataloader = DataLoader(train_dataset, batch_size=args.batch_size, sampler=sampler, pin_memory=use_gpu, num_workers=args.dataload_workers_nums) valid_dataloader = DataLoader(valid_dataset, batch_size=args.batch_size, shuffle=False, pin_memory=use_gpu, num_workers=args.dataload_workers_nums) # a name used to save checkpoints etc. # 2. prepare the model, checkpoint full_name = '%s_%s_%s_bs%d_lr%.1e_wd%.1e' % ( args.model, args.optim, args.lr_scheduler, args.batch_size, args.learning_rate, args.weight_decay) if args.comment: full_name = '%s_%s' % (full_name, args.comment) model = models.create_model(model_name=args.model, num_classes=len(CLASSES), in_channels=1) if use_gpu: model = torch.nn.DataParallel(model).cuda() criterion = torch.nn.CrossEntropyLoss() if args.optim == 'sgd': optimizer = torch.optim.SGD(model.parameters(), lr=args.learning_rate, momentum=0.9, weight_decay=args.weight_decay) else: optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay) start_timestamp = int(time.time() * 1000) start_epoch = 0 best_accuracy = 0 best_loss = 1e100 global_step = 0 if args.resume: print("resuming getShapeLista checkpoint '%s'" % args.resume) checkpoint = torch.load(args.resume) model.load_state_dict(checkpoint['state_dict']) model.float() optimizer.load_state_dict(checkpoint['optimizer']) best_accuracy = checkpoint.get('accuracy', best_accuracy) best_loss = checkpoint.get('loss', best_loss) start_epoch = checkpoint.get('epoch', start_epoch) global_step = checkpoint.get('step', global_step) del checkpoint # reduce memory if args.lr_scheduler == 'plateau': lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer, patience=args.lr_scheduler_patience, factor=args.lr_scheduler_gamma) else: lr_scheduler = torch.optim.lr_scheduler.StepLR( optimizer, step_size=args.lr_scheduler_step_size, gamma=args.lr_scheduler_gamma, last_epoch=start_epoch - 1) def get_lr(): return optimizer.param_groups[0]['lr'] writer = SummaryWriter(comment=('_speech_commands_' + full_name)) #3. train and validation print("training %s for Google speech commands..." % args.model) since = time.time() #grad_client_list = [[]] * args.clients global_grad_sum = np.zeros((84454500)) for epoch in range(start_epoch, args.max_epochs): print("epoch %3d with lr=%.02e" % (epoch, get_lr())) phase = 'train' writer.add_scalar('%s/learning_rate' % phase, get_lr(), epoch) model.train() # Set model to training mode running_loss = 0.0 it = 0 correct = 0 total = 0 A = 0 #federated n = 0 # the length of squeezed gradient vector shape_list = [] num_paddings = 0 A_inv = 0 S_i = 0 S_j = 0 u = 0 sigma = 0 vh = 0 ns = 0 #compute for each client current_client = 0 pbar = tqdm(train_dataloader, unit="audios", unit_scale=train_dataloader.batch_size, disable=True) for batch in pbar: inputs = batch['input'] inputs = torch.unsqueeze(inputs, 1) targets = batch['target'] #print(inputs.shape, targets.shape) if args.mixup: inputs, targets = mixup(inputs, targets, num_classes=len(CLASSES)) inputs = Variable(inputs, requires_grad=True) targets = Variable(targets, requires_grad=False) if use_gpu: inputs = inputs.cuda() targets = targets.cuda(async=True) outputs = model(inputs) if args.mixup: loss = mixup_cross_entropy_loss(outputs, targets) else: loss = criterion(outputs, targets) optimizer.zero_grad() loss.backward() #generate gradient list current_client_grad = [] for name, param in model.named_parameters(): if param.requires_grad: #print(name, param.grad.shape, param.grad.type())#, param.grad) current_client_grad.append(param.grad) #break #print(len(current_client_grad), current_client_grad[0].shape, current_client_grad[-1].shape) #randomize the gradient, if in a new batch, generate the randomization matrix if (current_client == 0): # Generate matrix A # 1. first obtain n(the total length of gradient vector) (if n == 0, get it, or pass) n = getLenOfGradientVector(current_client_grad) shape_list = getShapeList(current_client_grad) print("gradient vector of length", n) # 2. randomize a full rank matrix A, the elements are evenly distributed #A = np.random.randint(0, 10000,size = (n, n)) ## TODO Memory is enough to create such a large matrix : CURRENT SOLUTION use a small size matrix and # iterate over the vector A = np.random.randint(0, 10000, size=(args.matrix_size, args.matrix_size)) print("generating randomization matrix") A_inv = np.linalg.inv(A) # two index set S_i = random.sample(range(0, 3 * args.matrix_size), args.matrix_size) S_i.sort() S_j = random.sample(range(0, 2 * args.matrix_size), args.matrix_size) S_j.sort() # extend to B B = np.zeros((args.matrix_size, 3 * args.matrix_size)) for j in range(0, args.matrix_size): for i in range(0, args.matrix_size): B[i][S_i[j]] = A[i][j] #C = np.zeros((2 * args.matrix_size, 3*args.matrix_size)) C = np.random.randint(0, 10000, size=(2 * args.matrix_size, 3 * args.matrix_size)) for i in range(0, args.matrix_size): for j in range(0, 3 * args.matrix_size): C[S_j[i]][j] = B[i][j] # SVD u, s, vh = np.linalg.svd(C, full_matrices=True) print('u', u.shape, 's', s.shape, 'vh', vh.shape) # recontruct sigma sigma = np.zeros((C.shape[0], C.shape[1])) #print(C.shape[0], C.shape[1]) sigma[:min(C.shape[0], C.shape[1] ), :min(C.shape[0], C.shape[1])] = np.diag(s) assert (np.allclose(C, np.dot(u, np.dot(sigma, vh))) == True) # linear independent group(null space) ns = scipy.linalg.null_space( C) # (3000, 1000) we use the first args.matrix_size print('ns', ns.shape) # do the randomization, obtain a new gradient vector for gradient_client_list[current_client] flatterned_grad = transListOfArraysToArrays(current_client_grad, n) ###TODO: 1. need padding, 2. how to recover # random numbers r = np.random.randint(0, 10000, size=(args.matrix_size, 1)) r_new = np.zeros((3 * args.matrix_size, 1)) for i in range(args.matrix_size): r_new += r[i] * ns[:, i:i + 1] num_paddings = math.ceil( float(flatterned_grad.shape[0]) / args.matrix_size) * args.matrix_size - flatterned_grad.shape[0] print('flatterned', flatterned_grad.shape, 'need padding', num_paddings) #np.array # extent to 2n flatterned_grad_extended = np.zeros( (flatterned_grad.shape[0] + num_paddings)) flatterned_grad_extended[:flatterned_grad. shape[0]] = flatterned_grad #print(flatterned_grad_extended[:20]) current_client_grad_after_random = np.zeros( (3 * flatterned_grad_extended.shape[0])) new_grad = np.random.randint(0, 10000, size=(3 * args.matrix_size, 1)) for i in range(0, flatterned_grad_extended.shape[0], args.matrix_size): if (i / args.matrix_size % 5000 == 0): print(i / args.matrix_size, flatterned_grad_extended.shape[0] / args.matrix_size) for j in range(args.matrix_size): new_grad[S_i[j]] = flatterned_grad_extended[i + j] # compute the randomize gradient #print((new_grad + r_new).shape) randomized_gradient = np.dot(vh, new_grad + r_new) #print("randomized gradient", randomized_gradient.shape) current_client_grad_after_random[ 3 * i:3 * i + 3 * args.matrix_size] = np.squeeze(randomized_gradient) print('after randomization', current_client_grad_after_random.shape) # transform the flatterned vector to matrix for model update if (current_client == args.clients - 1): print("client", current_client) global_grad_sum += current_client_grad_after_random # collect all the randomized gradient, cacluate the sum, and send to all clients # each client eliminate the randomness, and update the parameters according to the gradients # remove randomness res = np.zeros((int(global_grad_sum.shape[0] / 3))) alpha = np.zeros((args.matrix_size, 1)) for i in range(0, global_grad_sum.shape[0], 3 * args.matrix_size): tmp = np.dot( u, np.dot(sigma, global_grad_sum[i:i + 3 * args.matrix_size])) for j in range(args.matrix_size): alpha[j] = tmp[S_j[j]] res[int(i / 3):int(i / 3) + args.matrix_size] = (np.dot( A_inv, alpha)).squeeze() # set the gradient manually and update recovered_grad_in_cuda = transNumpyGrad2Cuda(res, shape_list) ind = 0 #print(recovered_grad_in_cuda, recovered_grad_in_cuda[0].shape, r) for name, param in model.named_parameters(): if param.requires_grad: #print(recovered_grad_in_cuda[ind].shape, recovered_grad_in_cuda[ind].type()) #print(recovered_grad_in_cuda[ind][:20]) param.grad = recovered_grad_in_cuda[ind] ind += 1 assert (ind == len(recovered_grad_in_cuda)) optimizer.step() print("all clients finished") current_client = 0 global_grad_sum = np.zeros((84454500)) else: print("client", current_client) global_grad_sum += current_client_grad_after_random current_client += 1 # only update the parameters when current_client == args.clients - 1 # statistics it += 1 global_step += 1 #running_loss += loss.data[0] running_loss += loss.item() pred = outputs.data.max(1, keepdim=True)[1] if args.mixup: targets = batch['target'] targets = Variable(targets, requires_grad=False).cuda(async=True) correct += pred.eq(targets.data.view_as(pred)).sum() total += targets.size(0) writer.add_scalar('%s/loss' % phase, loss.item(), global_step) # update the progress bar pbar.set_postfix({ 'loss': "%.05f" % (running_loss / it), 'acc': "%.02f%%" % (100 * float(correct) / total) }) print("loss\t ", running_loss / it, "\t acc \t", 100 * float(correct) / total) #break accuracy = float(correct) / total epoch_loss = running_loss / it writer.add_scalar('%s/accuracy' % phase, 100 * accuracy, epoch) writer.add_scalar('%s/epoch_loss' % phase, epoch_loss, epoch)