Add files via upload

diagnose.py

@@ -0,0 +1,82 @@
+import os
+import time
+import numpy as np
+import argparse
+import json
+import torch
+
+from src.utils import load_net, load_data, \
+                      get_sharpness, get_nonuniformity, \
+                      eval_accuracy, eval_output,\
+                      get_gradW, get_gradx
+
+def get_args():
+    argparser = argparse.ArgumentParser(description=__doc__)
+    argparser.add_argument('--gpuid',default='0,')
+    argparser.add_argument('--dataset',default='fashionmnist',
+                            help='dataset choosed, [fashionmnist] | cifar10, 1dfunction')
+    argparser.add_argument('--network', default='vgg')
+    argparser.add_argument('--num_classes', type=int, default=2)
+    argparser.add_argument('--n_samples_per_class', type=int,
+                           default=500, help='training set size, [1000]')
+    argparser.add_argument('--batch_size', type=int,
+                            default=1000, help='batch size')
+    argparser.add_argument('--k', type=int, default=1)
+    argparser.add_argument('--model_file', default='fnn.pkl',
+                            help='file name of the pretrained model')
+    argparser.add_argument('--res_file', default='fnn_res.npz')
+    args = argparser.parse_args()
+    os.environ["CUDA_VISIBLE_DEVICES"]=args.gpuid
+
+    print('===> Config:')
+    print(json.dumps(vars(args),indent=2))
+    return args
+
+
+def main():
+    args = get_args()
+
+    # load model
+    criterion = torch.nn.MSELoss().cuda()
+    train_loader,test_loader = load_data(args.dataset,
+                                         args.num_classes,
+                                         train_per_class=
+                                             args.n_samples_per_class,
+                                         batch_size=1)
+    net = load_net(args.network, args.dataset, args.num_classes)
+    net.load_state_dict(torch.load('res/'+args.model_file))
+
+    # Evaluate models
+    train_loss, train_accuracy = eval_accuracy(net, criterion, train_loader)
+    test_loss, test_accuracy = eval_accuracy(net, criterion, test_loader)
+
+    print('===> Basic information of the given model: ')
+    print('\t train loss: %.2e, acc: %.2f'%(train_loss, train_accuracy))
+    print('\t test loss: %.2e, acc: %.2f'%(test_loss, test_accuracy))
+
+    print('===> Calculating output: ')
+    train_X, train_y = eval_output(net, train_loader)
+    test_X, test_y = eval_output(net, test_loader)
+
+    print('===> Compute gradient W:')
+    gradW = get_gradW(net, train_loader, ndata=args.num_classes*args.n_samples_per_class, k=args.k)
+    print('gradient W is %.3e\n' % (gradW))
+
+    print('===> Compute gradient x:')
+    gradx = get_gradx(net, train_loader, ndata=args.num_classes*args.n_samples_per_class, k=args.k)
+    print('gradient x is %.3e\n' % (gradx))
+
+    # print('===> Compute sharpness:')
+    # sharpness = get_sharpness(net, criterion, train_loader, \
+    #                             n_iters=10, verbose=True, tol=1e-4)
+    # print('Sharpness is %.2e\n'%(sharpness))
+
+    # print('===> Compute non-uniformity:')
+    # non_uniformity = get_nonuniformity(net, criterion, train_loader, \
+    #                                     n_iters=10, verbose=True, tol=1e-4)
+    # print('Non-uniformity is %.2e\n'%(non_uniformity))
+
+    np.savez('res/'+args.res_file, train_loss=train_loss, train_acc=train_accuracy, test_loss=test_loss, test_acc=test_accuracy, gradW=gradW, gradx=gradx, train_X=train_X, train_y=train_y, test_X=test_X, test_y=test_y)
+
+if __name__ == '__main__':
+    main()

postprocessing.py

@@ -0,0 +1,47 @@
+import numpy as np
+import matplotlib.pyplot as plt
+# bs = [10, 15, 20, 25, 30, 35, 40, 45, 50]
+lr = [0.02, 0.1, 0.5]
+rep = 5
+prefix = 'res/1d/res'
+
+N = len(lr) * rep
+test_error = np.zeros((N))
+gradW = np.zeros((N))
+gradx = np.zeros((N))
+#batch_size = np.zeros((N))
+learning_rate = np.zeros((N))
+train_X, train_y = [], []
+test_X, test_y = [], []
+
+i = 0
+j = 0
+for l in lr:
+    for r in range(rep):
+        res = np.load(prefix+'_lr'+str(j+1)+'_'+str(r+1)+'.npz')
+        test_error[i] = res['test_loss']
+        gradW[i] = res['gradW']
+        gradx[i] = res['gradx']
+        #batch_size[i] = b
+        learning_rate[i] = l
+
+        train_X.append(res['train_X'])
+        train_y.append(res['train_y'])
+        test_X.append(res['test_X'])
+        test_y.append(res['test_y'])
+        i += 1
+    j += 1
+
+ax = plt.scatter(gradW, gradx, c=learning_rate)
+cb = plt.colorbar(ax)
+cb.set_label('learning rate', fontsize=16)
+plt.xlabel(r'$g_w$', fontsize=16)
+plt.ylabel(r'$g_x$', fontsize=16)
+# plt.savefig('figs/cifar_bs.png', bbox_inches='tight')
+plt.show()
+
+# plt.plot(gradx, test_error, '.')
+# plt.show()
+
+# plt.plot(gradW, test_error, '.')
+# plt.show()

run.sh

@@ -0,0 +1,17 @@
+#!/bin/bash
+#learning_rate=0.1
+batch_size=2
+
+i=1
+for learning_rate in 0.02 0.1 0.5
+do
+  for rep in 1 2 3 4 5
+  do
+    python train.py --dataset 1dfunction --network fnn --n_samples_per_class 10 --n_iters 100000 --batch_size $batch_size --load_size $batch_size --learning_rate $learning_rate --model_file 1d/model_lr${i}_${rep}.pkl
+  done
+  i=$((i+1))
+done
+
+
+
+

train.py

@@ -0,0 +1,81 @@
+import os
+import argparse
+import json
+import torch
+
+from src.trainer import train
+from src.utils import load_net, load_data, eval_accuracy
+
+
+def get_args():
+    argparser = argparse.ArgumentParser(description=__doc__)
+    argparser.add_argument('--gpuid',
+                          default='0,', help='gpu id, [0] ')
+    argparser.add_argument('--dataset', 
+                          default='fashionmnist', help='dataset, [fashionmnist] | cifar10, 1dfunction')
+    argparser.add_argument('--network', default='vgg',
+                           help='network, [vgg] | fnn, resnet')
+    argparser.add_argument('--num_classes', type=int, default=2)
+    argparser.add_argument('--n_samples_per_class', type=int,
+                           default=500, help='training set size, [1000]')
+    argparser.add_argument('--load_size', type=int,
+                           default=1000, help='load size for dataset, [1000]')
+    argparser.add_argument('--optimizer', 
+                           default='sgd', help='optimizer, [sgd]')
+    argparser.add_argument('--n_iters', type=int,
+                           default=10000, help='number of iteration used to train nets, [10000]')
+    argparser.add_argument('--batch_size', type=int,
+                           default=1000, help='batch size, [1000]')
+    argparser.add_argument('--learning_rate', type=float,
+                           default=1e-1, help='learning rate')
+    argparser.add_argument('--momentum', type=float,
+                           default='0.0', help='momentum, [0.0]')
+    argparser.add_argument('--model_file', 
+                           default='fnn.pkl', help='filename to save the net, fnn.pkl')
+
+    args = argparser.parse_args()
+    if args.load_size > args.batch_size:
+        raise ValueError('load size should not be larger than batch size')
+    # os.environ["CUDA_VISIBLE_DEVICES"] = args.gpuid
+
+    print('===> Config:')
+    print(json.dumps(vars(args), indent=2))
+    return args
+
+def get_optimizer(net, args):
+    if args.optimizer == 'sgd':
+        return torch.optim.SGD(net.parameters(), lr=args.learning_rate, momentum=args.momentum)
+    elif args.optimizer == 'adam':
+        return torch.optim.Adam(net.parameters(), lr=args.learning_rate)
+    else:
+        raise ValueError('optimizer %s has not been supported'%(args.optimizer))
+
+def main():
+    args = get_args()
+
+    criterion = torch.nn.MSELoss()
+    train_loader, test_loader = load_data(args.dataset,
+                                          args.num_classes,
+                                          train_per_class=args.n_samples_per_class,
+                                          batch_size=args.load_size)
+    net = load_net(args.network, args.dataset, args.num_classes)
+    optimizer = get_optimizer(net, args)
+    print(optimizer)
+
+    print('===> Architecture:')
+    print(net)
+
+    print('===> Start training')
+    train(net, criterion, optimizer, train_loader, args.batch_size, args.n_iters, verbose=True)
+
+    train_loss, train_accuracy = eval_accuracy(net, criterion, train_loader)
+    test_loss, test_accuracy = eval_accuracy(net, criterion, test_loader)
+    print('===> Solution: ')
+    print('\t train loss: %.2e, acc: %.2f' % (train_loss, train_accuracy))
+    print('\t test loss: %.2e, acc: %.2f' % (test_loss, test_accuracy))
+
+    torch.save(net.state_dict(), 'res/'+args.model_file)
+
+
+if __name__ == '__main__':
+    main()