Ozzafar/classification sophon

classification/finetune_test.py

@@ -37,49 +37,31 @@ def args_parser():
 import wandb
 import timm
 
-
-def test_finetune(model, trainset, testset, epochs, lr):
-    model = nn.DataParallel(model)
-    trainloader = DataLoader(trainset, batch_size=256, shuffle=True, num_workers=4,drop_last=True)
-    testloader = DataLoader(testset, batch_size=256, shuffle=False, num_workers=4,drop_last=True)
-    optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=1e-4)
-    criterion = nn.CrossEntropyLoss()
-    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)
-    model.train()
-    # epochs = 1
-    for ep in tqdm(range(epochs)):
-        for inputs, targets in tqdm(trainloader):
-            inputs, targets = inputs.cuda(), targets.cuda()  
-            outputs = model(inputs)
-            loss = criterion(outputs, targets)
-            optimizer.zero_grad()
-            loss.backward()
-            optimizer.step()
-        # scheduler.step()
-    model.eval()
-    acc, test_loss = test(model, testloader, torch.device('cuda'))
-    return round(acc,2), round(test_loss,2)
-
 def test_finetune_final(args, mode, model, trainset, testset, epochs, lr):
     model = nn.DataParallel(model)
     trainloader = DataLoader(trainset, batch_size=args.bs, shuffle=True, num_workers=4,drop_last=True)
     testloader = DataLoader(testset, batch_size=args.bs, shuffle=False, num_workers=4,drop_last=True)
-    optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=1e-4)
+    # optimizer = optim.SGD(model.module.score.parameters(), lr=lr, momentum=0.9, weight_decay=1e-4)
+    optimizer = optim.SGD(model.module.score.parameters(), lr=lr, weight_decay=0.01)
+
+
     criterion = nn.CrossEntropyLoss()
-    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)
+    # scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)
     model.train()
     accs = []
     losses = []
-    # epochs = 1
-    for ep in tqdm(range(epochs)):
+    for ep in tqdm(range(epochs), desc="epoch"):
         model.train()
-        for inputs, targets in trainloader:
-            inputs, targets = inputs.cuda(), targets.cuda()  
-            outputs = model(inputs)
-            loss = criterion(outputs, targets)
+        train_loss = 0
+        for batch in tqdm(trainloader, desc="batch"):
+            inputs, targets, attention_mask = torch.stack(batch["input_ids"], dim=1).cuda(), batch["label"].cuda(), torch.stack(batch["attention_mask"], dim=1).cuda()
+            outputs = model(inputs, attention_mask=attention_mask)
+            loss = criterion(outputs.logits, targets)
+            train_loss += loss.item()
             optimizer.zero_grad()
             loss.backward()
             optimizer.step()
+        wandb.log({f'{mode}: train loss': train_loss / len(trainloader.dataset)})
         test_acc, test_loss = test(model, testloader, torch.device('cuda'))
         accs.append(test_acc)
         losses.append(test_loss)
@@ -89,13 +71,12 @@ def test_finetune_final(args, mode, model, trainset, testset, epochs, lr):
 
 if __name__ == '__main__':
     args = args_parser()
-    import wandb 
+    import wandb
     wandb.init(
-    project="sohpon classification finetune test",  
-    entity="sophon",
-    config = args,
-    name = f"{args.arch}_{args.dataset}" ,
-    notes = args.notes)   
+        project="sophon classification finetune test nlp",
+        config = args,
+        name = f"{args.arch}_{args.dataset}" ,
+        notes = args.notes)
     seed = args.seed
     set_seed(seed)
     trainset_tar, testset_tar = get_dataset(args.dataset, '../../../datasets', args=args)
@@ -113,9 +94,11 @@ def test_finetune_final(args, mode, model, trainset, testset, epochs, lr):
         acc, test_loss = test_finetune_final(args, 'normal pretrained/direct all', model.cuda(), trainset_tar, testset_tar, args.truly_finetune_epochs, args.finetune_lr)
 
     # ### train from scratch
-    elif args.start == 'sratch':
+    elif args.start == 'scratch':
         print('========test train from scratch=========')
         acc, test_loss = test_finetune_final(args, 'train from scratch/', model.cuda(), trainset_tar, testset_tar, args.truly_finetune_epochs, args.finetune_lr)
-    
+
     else:
-        assert(0)
+        assert(0)
+
+    print(f'test accuracy is {acc}, test loss is {test_loss}')

classification/inverse_loss.py

@@ -6,6 +6,8 @@
 import argparse
 import json
 import sys
+from transformers import GPT2LMHeadModel
+
 sys.path.append('../')
 def args_parser():
     parser = argparse.ArgumentParser(description='train N shadow models')
@@ -19,7 +21,7 @@ def args_parser():
     parser.add_argument('--test_iterval', default=10, type=int)
     parser.add_argument('--arch', default='caformer', type=str)
     parser.add_argument('--gpus', default='0,1', type=str)
-    parser.add_argument('--dataset', default='', type=str, choices=['CIFAR10', 'MNIST', 'SVHN', 'STL', 'CINIC'])
+    parser.add_argument('--dataset', default='', type=str, choices=['CIFAR10', 'MNIST', 'SVHN', 'STL', 'CINIC', 'IMDB', 'PILE'])
     parser.add_argument('--finetune_epochs', default=1, type=int)
     parser.add_argument('--truly_finetune_epochs', default=20, type=int)
     parser.add_argument('--finetune_lr', default=0.0001, type=float)
@@ -52,9 +54,9 @@ def fast_adapt_multibatch(batches, learner, loss, shots, ways, device):
     test_loss = 0
     test_accuracy = 0
     total_test = 0
+
     for index,batch in enumerate(batches):
-        data, labels = batch
-        data, labels = data.to(device), labels.to(device)
+        data, labels, attention_mask = torch.stack(batch["input_ids"], dim=1).cuda(), batch["label"].cuda(), torch.stack(batch["attention_mask"], dim=1).cuda()
         adaptation_indices = np.zeros(data.size(0), dtype=bool)
         # adaptation_indices[np.arange(shots*ways)] = True
         adaptation_indices[np.random.choice(np.arange(data.size(0)), shots*ways, replace=False)] = True
@@ -65,13 +67,13 @@ def fast_adapt_multibatch(batches, learner, loss, shots, ways, device):
         current_test = evaluation_data.shape[0]
         # print(current_test)
         total_test += current_test
-        adaptation_error = loss(learner(adaptation_data), adaptation_labels)
+        adaptation_error = loss(learner(adaptation_data).logits, adaptation_labels)
         if index == 0:
-            current_grads = learner.adapt(adaptation_error,None) 
+            current_grads = learner.adapt(adaptation_error,None, allow_nograd=True) #allow_nograd?
         else:
             last_grads = current_grads
-            current_grads = learner.adapt(adaptation_error,last_grads) 
-        predictions = learner(evaluation_data)
+            current_grads = learner.adapt(adaptation_error,last_grads, allow_nograd=True) 
+        predictions = learner(evaluation_data).logits
         evaluation_error = loss(1-predictions, evaluation_labels)  
         evaluation_accuracy = accuracy(predictions, evaluation_labels)
         test_loss += evaluation_error*current_test
@@ -81,9 +83,9 @@ def fast_adapt_multibatch(batches, learner, loss, shots, ways, device):
 
 def test_finetune(model, trainset, testset, epochs, lr):
     model = nn.DataParallel(model)
-    trainloader = DataLoader(trainset, batch_size=256, shuffle=True, num_workers=4,drop_last=True)
-    testloader = DataLoader(testset, batch_size=256, shuffle=False, num_workers=4,drop_last=True)
-    optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=1e-4)
+    trainloader = DataLoader(trainset, batch_size=8, shuffle=True, num_workers=4,drop_last=True)
+    testloader = DataLoader(testset, batch_size=8, shuffle=False, num_workers=4,drop_last=True)
+    optimizer = optim.SGD(model.parameters(), lr=lr, weight_decay=0.01)
     criterion = nn.CrossEntropyLoss()
     scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)
     model.train()
@@ -101,19 +103,19 @@ def test_finetune(model, trainset, testset, epochs, lr):
 
 def test_finetune_final(mode, model, trainset, testset, epochs, lr):
     model = nn.DataParallel(model)
-    trainloader = DataLoader(trainset, batch_size=256, shuffle=True, num_workers=4,drop_last=True)
-    testloader = DataLoader(testset, batch_size=256, shuffle=False, num_workers=4,drop_last=True)
-    optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=1e-4)
+    trainloader = DataLoader(trainset, batch_size=8, shuffle=True, num_workers=4, drop_last=True)
+    testloader = DataLoader(testset, batch_size=8, shuffle=False, num_workers=4, drop_last=True)
+    optimizer = optim.SGD(model.parameters(), lr=lr, weight_decay=0.01)
     criterion = nn.CrossEntropyLoss()
     scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)
     model.train()
     # epochs = 1
     for ep in tqdm(range(epochs)):
         model.train()
-        for inputs, targets in tqdm(trainloader):
-            inputs, targets = inputs.cuda(), targets.cuda()  
-            outputs = model(inputs)
-            loss = criterion(outputs, targets)
+        for batch in tqdm(trainloader):
+            inputs, targets, attention_mask = torch.stack(batch["input_ids"], dim=1).cuda(), batch["label"].cuda(), torch.stack(batch["attention_mask"], dim=1).cuda()
+            outputs = model(inputs, attention_mask=attention_mask)
+            loss = criterion(outputs.logits, targets)
             optimizer.zero_grad()
             loss.backward()
             optimizer.step()
@@ -140,10 +142,11 @@ def main(
     print("Hostname:", hostname)
     ip_address = socket.gethostbyname(hostname)
     args.from_machine = ip_address
+    lm_model = GPT2LMHeadModel.from_pretrained("gpt2").cuda()
 
     wandb.init(
     project="sophon classification",  
-    entity="sophon",
+    # entity="sophon",
     config = args,
     name = f"{args.dataset}_alpha{args.alpha}_beta{args.beta}_ml{args.ml_loop}_nl{args.nl_loop}_batches{args.adaptation_steps}" ,
     notes= args.notes,         
@@ -157,14 +160,14 @@ def main(
         # torch.cuda.manual_seed(seed)
         device = torch.device('cuda')
     wandb.log({'seed':seed})
-        save_path = args.root + '/inverse_loss'+ '/'+args.arch+'_'+ args.dataset + '/'
+    save_path = args.root + '/inverse_loss'+ '/'+args.arch+'_'+ args.dataset + '/'
     adaptation_steps = args.adaptation_steps
     now = datetime.now()
     save_path = save_path + '/' + f'{now.month}_{now.day}_{now.hour}_{now.minute}_{now.second}/'
     os.makedirs(save_path, exist_ok=True)
     wandb.log({'save path': save_path})
     save_args_to_file(args, save_path+"args.json")
-    trainset_ori, testset_ori = get_dataset('ImageNet', '../../../datasets/', subset='imagenette', args=args)
+    trainset_ori, testset_ori = get_dataset("PILE", '../../../datasets', args=args) #get_dataset('ImageNet', '../../../datasets/', subset='imagenette', args=args)
     original_trainloader = DataLoader(trainset_ori, batch_size=args.bs, shuffle=True, num_workers=0)
     original_testloader = DataLoader(testset_ori, batch_size=args.bs, shuffle=False, num_workers=0)
     trainset_tar, testset_tar = get_dataset(args.dataset, '../../../datasets', args=args)
@@ -264,12 +267,19 @@ def main(
             except StopIteration:
                 original_iter = iter(original_trainloader)
                 batch = next(original_iter)
-            inputs, targets = batch
-            inputs, targets = inputs.cuda(), targets.cuda()       
-            # print(inputs.shape)
+            inputs, targets, attention_mask = torch.stack(batch["input_ids"], dim=1).cuda(), batch["label"].cuda(), torch.stack(batch["attention_mask"], dim=1).cuda()
             natural_optimizer.zero_grad()
-            outputs = model(inputs)
-            loss = criterion(outputs, targets) 
+            outputs = model(input_ids=inputs, attention_mask=attention_mask, output_hidden_states=True)
+
+            last_hidden_state = outputs.hidden_states[-1]
+            with torch.no_grad():
+                mask = attention_mask == 1
+                next_token_indexes = (mask.cumsum(dim=1) * mask).argmax(dim=1)
+                next_token_indexes[next_token_indexes == 0] = -1
+            last_token_hidden_state = last_hidden_state[range(last_hidden_state.shape[0]), next_token_indexes]
+            logits = lm_model.lm_head(last_token_hidden_state)
+
+            loss = criterion(logits, targets)
             loss.backward()
             avg_gradients = check_gradients(model)
             # print('check gradients!!!!!!!!!')

classification/kl_uniform_loss.py

@@ -21,7 +21,7 @@ def args_parser():
     parser.add_argument('--test_iterval', default=10, type=int)
     parser.add_argument('--arch', default='', type=str)
     parser.add_argument('--gpus', default='0,1', type=str)
-    parser.add_argument('--dataset', default='CIFAR10', type=str, choices=['CIFAR10', 'MNIST', 'SVHN', 'STL', 'CINIC'])
+    parser.add_argument('--dataset', default='CIFAR10', type=str, choices=['CIFAR10', 'MNIST', 'SVHN', 'STL', 'CINIC', 'IMDB'])
     parser.add_argument('--finetune_epochs', default=1, type=int)
     parser.add_argument('--truly_finetune_epochs', default=20, type=int)
     parser.add_argument('--finetune_lr', default=0.0001, type=float)
@@ -133,9 +133,9 @@ def partial_fast_adapt_multibatch(batches, learner, loss, shots, ways, device):
 
 def test_finetune(model, trainset, testset, epochs, lr):
     model = nn.DataParallel(model)
-    trainloader = DataLoader(trainset, batch_size=256, shuffle=True, num_workers=4,drop_last=True)
-    testloader = DataLoader(testset, batch_size=256, shuffle=False, num_workers=4,drop_last=True)
-    optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=1e-4)
+    trainloader = DataLoader(trainset, batch_size=8, shuffle=True, num_workers=4,drop_last=True)
+    testloader = DataLoader(testset, batch_size=8, shuffle=False, num_workers=4,drop_last=True)
+    optimizer = optim.SGD(model.parameters(), lr=lr, weight_decay=0.01)
     criterion = nn.CrossEntropyLoss()
     scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)
     model.train()
@@ -153,9 +153,9 @@ def test_finetune(model, trainset, testset, epochs, lr):
 
 def test_finetune_final(mode, model, trainset, testset, epochs, lr):
     model = nn.DataParallel(model)
-    trainloader = DataLoader(trainset, batch_size=256, shuffle=True, num_workers=4,drop_last=True)
-    testloader = DataLoader(testset, batch_size=256, shuffle=False, num_workers=4,drop_last=True)
-    optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=1e-4)
+    trainloader = DataLoader(trainset, batch_size=8, shuffle=True, num_workers=4,drop_last=True)
+    testloader = DataLoader(testset, batch_size=8, shuffle=False, num_workers=4,drop_last=True)
+    optimizer = optim.SGD(model.parameters(), lr=lr, weight_decay=0.01)
     criterion = nn.CrossEntropyLoss()
     scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)
     model.train()

classification/model.py

@@ -3,8 +3,12 @@
 import torch.nn as nn
 import torch.nn.functional as F
 import torch.utils.model_zoo as model_zoo
+from scipy.special import softmax
+from torch.nn.modules.module import T
 from torchvision.models.resnet import ResNet
 from torchvision.models.resnet import BasicBlock, Bottleneck
+from transformers import GPT2LMHeadModel, GPT2ForSequenceClassification, GPT2Tokenizer
+
 
 # 定义ResNet-18结构
 def resnet18(pretrained=False, **kwargs):
@@ -136,4 +140,14 @@ def vgg19(pretrained=False, **kwargs):
     'vgg13_bn': 'https://download.pytorch.org/models/vgg13_bn-abd245e5.pth',
     'vgg16_bn': 'https://download.pytorch.org/models/vgg16_bn-6c64b313.pth',
     'vgg19_bn': 'https://download.pytorch.org/models/vgg19_bn-c79401a0.pth',
-}
+}
+
+def gpt2(pretrained=False, **kwargs):
+    model = GPT2ForSequenceClassification.from_pretrained("gpt2", num_labels=2, **kwargs)
+    model.config.num_labels = 2
+
+    tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
+    tokenizer.pad_token = tokenizer.eos_token  # Use eos_token as padding token
+
+    model.config.pad_token_id = tokenizer.pad_token_id
+    return model