def test_get_prototypes(): # Numpy num_classes = 3 embeddings_np = np.random.rand(2, 5, 7).astype(np.float32) targets_np = np.random.randint(0, num_classes, size=(2, 5)) # PyTorch embeddings_th = torch.as_tensor(embeddings_np) targets_th = torch.as_tensor(targets_np) prototypes_th = get_prototypes(embeddings_th, targets_th, num_classes) assert prototypes_th.shape == (2, num_classes, 7) assert prototypes_th.dtype == embeddings_th.dtype prototypes_np = np.zeros((2, num_classes, 7), dtype=np.float32) num_samples_np = np.zeros((2, num_classes), dtype=np.int_) for i in range(2): for j in range(5): num_samples_np[i, targets_np[i, j]] += 1 for k in range(7): prototypes_np[i, targets_np[i, j], k] += embeddings_np[i, j, k] for i in range(2): for j in range(num_classes): for k in range(7): prototypes_np[i, j, k] /= max(num_samples_np[i, j], 1) np.testing.assert_allclose(prototypes_th.detach().numpy(), prototypes_np)
def rep_memory(args, model, memory_train): memory_loss = 0 for dataidx, dataloader_dict in enumerate(memory_train): for dataname, memory_list in dataloader_dict.items(): select = random.choice(memory_list) memory_train_inputs, memory_train_targets = select['train'] memory_train_inputs = memory_train_inputs.to(device=args.device) memory_train_targets = memory_train_targets.to(device=args.device) if memory_train_inputs.size(2) == 1: memory_train_inputs = memory_train_inputs.repeat(1, 1, 3, 1, 1) memory_train_embeddings = model(memory_train_inputs, dataidx) memory_test_inputs, memory_test_targets = select['test'] memory_test_inputs = memory_test_inputs.to(device=args.device) memory_test_targets = memory_test_targets.to(device=args.device) if memory_test_inputs.size(2) == 1: memory_test_inputs = memory_test_inputs.repeat(1, 1, 3, 1, 1) memory_test_embeddings = model(memory_test_inputs, dataidx) memory_prototypes = get_prototypes(memory_train_embeddings, memory_train_targets, args.num_way) memory_loss += prototypical_loss(memory_prototypes, memory_test_embeddings, memory_test_targets) return memory_loss
def test(device, testset, testloader, model): model.to(device) model.eval() acc = [] with tqdm(testloader, total=config.num_batches) as pbar: for batch_idx, batch in enumerate(pbar): train_inputs, train_targets = batch['train'] train_inputs = train_inputs.to(device=device) train_targets = train_targets.to(device=device) train_embeddings = model(train_inputs) test_inputs, test_targets = batch['test'] test_inputs = test_inputs.to(device=device) test_targets = test_targets.to(device=device) test_embeddings = model(test_inputs) prototypes = get_prototypes(train_embeddings, train_targets, testset.num_classes_per_task) with torch.no_grad(): accuracy = get_accuracy(prototypes, test_embeddings, test_targets) pbar.set_postfix(accuracy='{0:.4f}'.format(accuracy.item())) acc.append(accuracy) if batch_idx >= config.num_batches: break return acc
def train(args): logger.warning('This script is an example to showcase the extensions and ' 'data-loading features of Torchmeta, and as such has been ' 'very lightly tested.') dataset = omniglot(args.folder, shots=args.num_shots, ways=args.num_ways, shuffle=True, test_shots=15, meta_train=True, download=args.download) dataloader = BatchMetaDataLoader(dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers) model = PrototypicalNetwork(1, args.embedding_size, hidden_size=args.hidden_size) model.to(device=args.device) model.train() optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) # Training loop with tqdm(dataloader, total=args.num_batches) as pbar: for batch_idx, batch in enumerate(pbar): model.zero_grad() train_inputs, train_targets = batch['train'] train_inputs = train_inputs.to(device=args.device) train_targets = train_targets.to(device=args.device) train_embeddings = model(train_inputs) test_inputs, test_targets = batch['test'] test_inputs = test_inputs.to(device=args.device) test_targets = test_targets.to(device=args.device) test_embeddings = model(test_inputs) prototypes = get_prototypes(train_embeddings, train_targets, dataset.num_classes_per_task) loss = prototypical_loss(prototypes, test_embeddings, test_targets) loss.backward() optimizer.step() with torch.no_grad(): accuracy = get_accuracy(prototypes, test_embeddings, test_targets) pbar.set_postfix(accuracy='{0:.4f}'.format(accuracy.item())) if batch_idx >= args.num_batches: break # Save model if args.output_folder is not None: filename = os.path.join(args.output_folder, 'protonet_omniglot_' '{0}shot_{1}way.pt'.format(args.num_shots, args.num_ways)) with open(filename, 'wb') as f: state_dict = model.state_dict() torch.save(state_dict, f)
def train(device, dataset, dataloader, model): print("in train") model = model.to(device) model.train() optimizer = torch.optim.Adam(model.parameters(), lr=config.learning_rate) # Training loop images_per_batch = {} batch_count, images_per_batch['train'], images_per_batch[ 'test'] = 0, [], [] with tqdm(dataloader, total=config.num_batches) as pbar: for batch_idx, batch in enumerate(pbar): model.zero_grad() train_inputs, train_targets = batch['train'] train_inputs = train_inputs.to(device=device) train_targets = train_targets.to(device=device) train_embeddings = model(train_inputs) test_inputs, test_targets = batch['test'] test_inputs = test_inputs.to(device=device) test_targets = test_targets.to(device=device) test_embeddings = model(test_inputs) prototypes = get_prototypes(train_embeddings, train_targets, dataset.num_classes_per_task) loss = prototypical_loss(prototypes, test_embeddings, test_targets) loss.backward() optimizer.step() #Just keeping the count here batch_count += 1 images_per_batch['train'].append(train_inputs.shape[1]) images_per_batch['test'].append(test_inputs.shape[1]) with torch.no_grad(): accuracy = get_accuracy(prototypes, test_embeddings, test_targets) pbar.set_postfix(accuracy='{0:.4f}'.format(accuracy.item())) if batch_idx >= config.num_batches: break print("Number of batches in the dataloader: ", batch_count) # Save model if check_dir() is not None: filename = os.path.join( 'saved_models', 'protonet_cifar_fs_{0}shot_{1}way.pt'.format(config.k, config.n)) with open(filename, 'wb') as f: state_dict = model.state_dict() torch.save(state_dict, f) print("Model saved") return batch_count, images_per_batch
def valid(self, dataloader_dict, domain_id, epoch): self.model.eval() acc_dict = {} acc_list = [] for dataname, dataloader in dataloader_dict.items(): with torch.no_grad(): with tqdm(dataloader, total=self.args.num_valid_batches) as pbar: for batch_idx, batch in enumerate(pbar): self.model.zero_grad() train_inputs, train_targets = batch['train'] train_inputs = train_inputs.to(device=self.args.device) train_targets = train_targets.to( device=self.args.device) if train_inputs.size(2) == 1: train_inputs = train_inputs.repeat(1, 1, 3, 1, 1) train_embeddings = self.model(train_inputs, domain_id) test_inputs, test_targets = batch['test'] test_inputs = test_inputs.to(device=self.args.device) test_targets = test_targets.to(device=self.args.device) if test_inputs.size(2) == 1: test_inputs = test_inputs.repeat(1, 1, 3, 1, 1) test_embeddings = self.model(test_inputs, domain_id) prototypes = get_prototypes(train_embeddings, train_targets, self.args.num_way) accuracy = get_accuracy(prototypes, test_embeddings, test_targets) acc_list.append(accuracy.cpu().data.numpy()) pbar.set_description( 'dataname {} accuracy ={:.4f}'.format( dataname, np.mean(acc_list))) if batch_idx >= self.args.num_valid_batches: break avg_accuracy = np.round(np.mean(acc_list), 4) acc_dict = {dataname: avg_accuracy} return acc_dict
def train(self, epoch, dataloader_dict, memory_train=None, domain_id=None): self.model.train() for dataname, dataloader in dataloader_dict.items(): with tqdm(dataloader, total=self.args.num_batches) as pbar: for batch_idx, batch in enumerate(pbar): self.model.zero_grad() train_inputs, train_targets = batch['train'] train_inputs = train_inputs.to(device=self.args.device) train_targets = train_targets.to(device=self.args.device) if train_inputs.size(2) == 1: train_inputs = train_inputs.repeat(1, 1, 3, 1, 1) train_embeddings = self.model(train_inputs, domain_id) test_inputs, test_targets = batch['test'] test_inputs = test_inputs.to(device=self.args.device) test_targets = test_targets.to(device=self.args.device) if test_inputs.size(2) == 1: test_inputs = test_inputs.repeat(1, 1, 3, 1, 1) test_embeddings = self.model(test_inputs, domain_id) prototypes = get_prototypes(train_embeddings, train_targets, args.num_way) loss = prototypical_loss(prototypes, test_embeddings, test_targets) loss.backward(retain_graph=True) param_list = [] param_names = [] for name, v in self.model.named_parameters(): if 'domain_out' not in name: if v.requires_grad: param_list.append(v) param_names.append(name) first_grad = torch.autograd.grad(loss, param_list, create_graph=False, retain_graph=False) val_graddict = {} layer_name = [] for gradient, name in zip(first_grad, param_names): split_name = name.split('.') layer = split_name[0] if layer not in self.args.layer_filters: if layer not in layer_name: layer_name.append(layer) val_graddict[layer] = [] val_graddict[layer].append( gradient.clone().view(-1)) else: val_graddict[layer].append( gradient.clone().view(-1)) else: layer_sub = layer + '.' + split_name[ 1] + '.' + split_name[2] if layer_sub not in layer_name: layer_name.append(layer_sub) val_graddict[layer_sub] = [] val_graddict[layer_sub].append( gradient.clone().view(-1)) else: val_graddict[layer_sub].append( gradient.clone().view(-1)) for key in val_graddict: val_graddict[key] = torch.cat(val_graddict[key]) self.optimizer.step() if memory_train: memory_trainnew = copy.deepcopy(memory_train) self.hyper_optim.optimizer = self.optimizer self.hyper_optim.compute_hg(self.model, val_graddict) val_grad = self.rep_grad_new(self.args, memory_trainnew) self.hyper_optim.hyper_step(val_grad) self.model.zero_grad() if batch_idx >= args.num_batches: break
def rep_grad_new(self, args, memory_train): memory_loss = 0 for dataidx, dataloader_dict in enumerate(memory_train): for dataname, memory_list in dataloader_dict.items(): select = random.choice(memory_list) memory_train_inputs, memory_train_targets = select['train'] memory_train_inputs = memory_train_inputs.to( device=args.device) memory_train_targets = memory_train_targets.to( device=args.device) if memory_train_inputs.size(2) == 1: memory_train_inputs = memory_train_inputs.repeat( 1, 1, 3, 1, 1) memory_train_embeddings = self.model(memory_train_inputs, dataidx) memory_test_inputs, memory_test_targets = select['test'] memory_test_inputs = memory_test_inputs.to(device=args.device) memory_test_targets = memory_test_targets.to( device=args.device) if memory_test_inputs.size(2) == 1: memory_test_inputs = memory_test_inputs.repeat( 1, 1, 3, 1, 1) indlist = [] for ind in range(len(memory_train)): if ind != dataidx: indlist.append(ind) if indlist: indselect = random.choice(indlist) dataloader_dict2 = memory_train[indselect] for dataname, memory_list in dataloader_dict2.items(): select2 = random.choice(memory_list) memory_train_inputs2, memory_train_targets2 = select2[ 'train'] memory_train_inputs2 = memory_train_inputs2.to( device=args.device) memory_train_targets2 = memory_train_targets2.to( device=args.device) if memory_train_inputs2.size(2) == 1: memory_train_inputs2 = memory_train_inputs2.repeat( 1, 1, 3, 1, 1) memory_train_embeddings2 = self.model( memory_train_inputs2, dataidx) memory_test_inputs2, memory_test_targets2 = select2['test'] memory_test_inputs2 = memory_test_inputs2.to( device=args.device) memory_test_targets2 = memory_test_targets2.to( device=args.device) if memory_test_inputs2.size(2) == 1: memory_test_inputs2 = memory_test_inputs2.repeat( 1, 1, 3, 1, 1) memory_test_embeddings2 = self.model(memory_test_inputs2, dataidx) memory_test_embeddings = self.model(memory_test_inputs, dataidx) memory_prototypes = get_prototypes(memory_train_embeddings, memory_train_targets, args.num_way) memory_loss += prototypical_loss(memory_prototypes, memory_test_embeddings, memory_test_targets) if indlist: sellist = range(len(memory_train_embeddings)) i = random.choice(sellist) memory_train_embeddings = torch.cat([ memory_train_embeddings[i].unsqueeze(0), memory_train_embeddings2[i].unsqueeze(0) ], dim=1) memory_train_targets = torch.cat([ memory_train_targets[i].unsqueeze(0), memory_train_targets2[i].unsqueeze(0) + args.num_way ], dim=1) memory_prototypes = get_prototypes(memory_train_embeddings, memory_train_targets, 2 * args.num_way) memory_test_embeddings = torch.cat([ memory_test_embeddings[i].unsqueeze(0), memory_test_embeddings2[i].unsqueeze(0) ], dim=1) memory_test_targets = torch.cat([ memory_test_targets[i].unsqueeze(0), memory_test_targets2[i].unsqueeze(0) + args.num_way ], dim=1) memory_loss += 1e-4 * prototypical_loss( memory_prototypes, memory_test_embeddings, memory_test_targets) param_list = [] param_names = [] for name, v in self.model.named_parameters(): if 'domain_out' not in name: if v.requires_grad: param_list.append(v) param_names.append(name) val_grad = torch.autograd.grad(memory_loss, param_list) val_graddict = {} layer_name = [] for gradient, name in zip(val_grad, param_names): split_name = name.split('.') layer = split_name[0] if layer not in self.args.layer_filters: if layer not in layer_name: layer_name.append(layer) val_graddict[layer] = [] val_graddict[layer].append(gradient.view(-1)) else: val_graddict[layer].append(gradient.view(-1)) else: layer_sub = layer + '.' + split_name[1] + '.' + split_name[2] if layer_sub not in layer_name: layer_name.append(layer_sub) val_graddict[layer_sub] = [] val_graddict[layer_sub].append(gradient.view(-1)) else: val_graddict[layer_sub].append(gradient.view(-1)) for key in val_graddict: val_graddict[key] = torch.cat(val_graddict[key]) self.model.zero_grad() memory_loss.detach_() return val_graddict
global_task_count += args.batch_tasks support_inputs, support_targets, support_semantics = get_inputs_and_outputs( args, train_batch['train']) query_inputs, query_targets, _ = get_inputs_and_outputs( args, train_batch['test']) support_embeddings, ca_weights, sa_weights = model( support_inputs, semantics=support_semantics, output_weights=True) query_embeddings = model(query_inputs) prototypes = get_prototypes(support_embeddings, support_targets, train_dataset.num_classes_per_task) train_loss = prototypical_loss(prototypes, query_embeddings, query_targets) train_acc = get_proto_accuracy(prototypes, query_embeddings, query_targets) del ca_weights, sa_weights optimizer.zero_grad() train_loss.backward() optimizer.step() pbar.set_postfix(train_acc='{0:.4f}'.format(train_acc.item())) # validation if global_task_count % args.valid_every_tasks == 0:
def train(args): datasets = [] for root_idx in tqdm(range(args.ds_size)): folders = { "index": root_idx, "cache": "/nrs/funke/wolfs2/lisl/datasets/prototypical_network_cache_uncompressed_fg2.zarr", "raw": ("/nrs/funke/wolfs2/lisl/datasets/dsb_indexed.zarr", f"train/raw/{root_idx}"), "gt_segmentation": ("/nrs/funke/wolfs2/lisl/datasets/dsb_indexed.zarr", f"train/gt_segmentation/{root_idx}"), "embedding": (("/nrs/funke/wolfs2/lisl/experiments/semantic/c32/prediction/anchor.zarr", f"train/prediction_interm/{root_idx}"), ("/nrs/funke/wolfs2/lisl/experiments/semantic/c32/prediction/semantic.zarr", f"train/prediction/{root_idx}")), "min_samples": 10, "bg_distance": 20 } ds = fast_dataset_creator(folders, shots=args.num_shots, ways=args.num_ways, shuffle=True, test_shots=3, transform=None, meta_train=True) if len(ds): datasets.append(ds) else: print( f"dataset with id {root_idx} appears to be empty. Will be skipped" ) os.makedirs(args.output_folder, exist_ok=True) loaders = [] for ds in datasets: loaders.append( BatchMetaDataLoader(ds, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers)) time.sleep(1) print("ds loaded") model = PrototypicalNetwork(544, args.inst_embedding_size, 2, hidden_size=args.hidden_size) model.to(device=args.device) model.train() optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) print("starting") accuracy_window = {} accuracy_window["inst_accuracy"] = [] accuracy_window["sem_accuracy"] = [] accuracy_window["combined"] = [] # Training loop for epoch in range(500): dataloader = roundrobin_break_early(*loaders) with tqdm(dataloader, total=args.num_batches) as pbar: for batch_idx, batch in enumerate(pbar): model.zero_grad() train_inputs, train_instance_targets, train_semantic_targets = batch[ 'train'] train_inputs = train_inputs.to(device=args.device) train_instance_targets = train_instance_targets.to( device=args.device) train_semantic_targets = train_semantic_targets.to( device=args.device) train_semantic_embeddings, train_spatial_instance_embeddings = model( train_inputs) test_inputs, test_instance_targets, test_semantic_targets = batch[ 'test'] test_inputs = test_inputs.to(device=args.device) test_instance_targets = test_instance_targets.to( device=args.device) test_semantic_targets = test_semantic_targets.to( device=args.device) test_semantic_embeddings, test_spatial_instance_embeddings = model( test_inputs) # semantic loss c = train_semantic_embeddings.shape[-1] semantic_loss = F.cross_entropy( train_semantic_embeddings.view(-1, c), train_semantic_targets.view(-1)) # semantic_prototypes = get_prototypes(train_semantic_embeddings, train_semantic_targets, 2) # semantic_loss = prototypical_loss(semantic_prototypes, test_semantic_embeddings, test_semantic_targets) # instance loss train_inst_emb = train_spatial_instance_embeddings test_inst_emb = test_spatial_instance_embeddings instance_prototypes = get_prototypes(train_inst_emb, train_instance_targets, args.num_ways) instance_loss = prototypical_loss(instance_prototypes, test_inst_emb, test_instance_targets) loss = semantic_loss + instance_loss loss.backward() optimizer.step() with torch.no_grad(): inst_accuracy = get_accuracy(instance_prototypes, test_inst_emb, test_instance_targets) pred = test_semantic_embeddings.max(-1).indices sem_accuracy = (pred == test_semantic_targets).sum().item() / ( pred.size(0) * pred.size(1)) sem_accuracy = float(sem_accuracy) if len(accuracy_window["inst_accuracy"]) > 100: accuracy_window["inst_accuracy"].pop(0) accuracy_window["sem_accuracy"].pop(0) accuracy_window["combined"].pop(0) accuracy_window["inst_accuracy"].append( float(inst_accuracy)) accuracy_window["sem_accuracy"].append(float(sem_accuracy)) accuracy_window["combined"].append( float(inst_accuracy + sem_accuracy) / 2) mean_inst_acc = np.mean(accuracy_window['inst_accuracy']) mean_sem_acc = np.mean(accuracy_window['sem_accuracy']) pbar.set_postfix(inst_accuracy=f"{mean_inst_acc:.4f}", sem_accuracy=f'{mean_sem_acc:.4f}', loss=f"{loss:.4}") if pbar.n > args.num_batches: break # Save model score = np.mean(accuracy_window["combined"]) if args.output_folder is not None: filename = os.path.join( args.output_folder, f'prototypical_networks_class_{epoch}_{args.num_shots}shot_{args.num_ways}way_{args.ds_size}_sem_spatial_{score:.4f}.pytorch' ) print("saving", filename) torch.save(model.state_dict(), filename)