if torch.cuda.is_available(): print('CUDA enabled.') net.cuda() print("--- Pretrained network loaded ---") # test(net, loader_test) # prune the weights masks = weight_prune(net, param['pruning_perc']) net.set_masks(masks) net = nn.DataParallel(net) print("--- {}% parameters pruned ---".format(param['pruning_perc'])) test(net, loader_test) # Retraining criterion = nn.CrossEntropyLoss() optimizer = torch.optim.RMSprop(net.parameters(), lr=param['learning_rate'], weight_decay=param['weight_decay']) train(net, criterion, optimizer, param, loader_train) # Check accuracy and nonzeros weights in each layer print("--- After retraining ---") test(net, loader_test) prune_rate(net) # Save and load the entire model torch.save(net.state_dict(), 'models/alexnet_pruned.pkl')
def main(): global args, best_prec1 args = parser.parse_args() print(args) # Set # classes if args.data == 'UCF101': num_classes = 101 else: num_classes = 0 print('Specify the dataset to use ') # Create model if args.pretrained: print("=> using pre-trained model '{}'".format(args.arch)) model = models.__dict__[args.arch](pretrained=True) else: print("=> creating model '{}'".format(args.arch)) model = models.__dict__[args.arch]() if args.arch.startswith('alexnet'): model = AlexNet(num_classes=num_classes) model.features = torch.nn.DataParallel(model.features) model.cuda() else: model = torch.nn.DataParallel(model).cuda() # Modify last layer of the model model_ft = models.resnet18(pretrained=True) num_ftrs = model_ft.fc.in_features model_ft.fc = nn.Linear(num_ftrs, 101) model = model_ft.cuda() model = torch.nn.DataParallel(model).cuda() # Using one GPU (device_ids = 1) # print(model) # Define loss function (criterion) and optimizer criterion = nn.CrossEntropyLoss().cuda() optimizer = torch.optim.SGD(model.parameters(), args.lr, momentum=args.momentum, weight_decay=args.weight_decay) # Optionally resume from a checkpoint if args.resume: checkpoint = torch.load(args.resume) args.start_epoch = checkpoint['epoch'] best_prec1 = checkpoint['best_prec1'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) print("=> loaded checkpoint '{}' (epoch {})" .format(args.resume, checkpoint['epoch'])) else: print("=> no checkpoint found at '{}'".format(args.resume)) cudnn.benchmark = True # Data loading code traindir = os.path.join(args.data, 'train') testdir = os.path.join(args.data, 'test') normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) train_loader = torch.utils.data.DataLoader( datasets.ImageFolder(traindir, transforms.Compose([ transforms.RandomCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize, ])), batch_size = args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True ) val_loader = torch.utils.data.DataLoader( datasets.ImageFolder(testdir, transforms.Compose([ transforms.Scale(256), transforms.CenterCrop(224), transforms.ToTensor(), normalize, ])), batch_size = args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True ) if args.evaluate: validate(val_loader, model, criterion) return for epoch in range(args.start_epoch, args.epochs): adjust_learning_rate(optimizer, epoch) # Train for one epoch train(train_loader, model, criterion, optimizer, epoch) # # Evaluate on validation set prec1 = validate(val_loader, model, criterion) # Remember best prec@1 and save checkpoint is_best = prec1 > best_prec1 best_prec1 = max(prec1, best_prec1) save_checkpoint({ 'epoch': epoch + 1, 'arch': args.arch, 'state_dict': model.state_dict(), 'best_prec1': best_prec1, 'optimizer': optimizer.state_dict(), }, is_best)
dtest = torchvision.datasets.CIFAR10( root='./data', train=False, download=True, transform=transform_test ) test_loader = torch.utils.data.DataLoader( dtest, batch_size=BATCH_SIZE, shuffle=False ) device = 'cuda' if torch.cuda.is_available() else 'cpu' # 損失関数にクロスエントロピー誤差を用いる criterion = nn.CrossEntropyLoss() # モデルの定義 model = AlexNet(num_classes=num_classes) # optimizerはSGD optimizer = torch.optim.SGD(model.parameters(), lr=LR, momentum=0.9, weight_decay=5e-4) # 実際の学習部分 for epoch in range(1, NUM_EPOCHS + 1): train_loss, train_acc = epoch_train(train_loader, model, optimizer, criterion) test_loss, test_acc = epoch_eval(test_loader, model, criterion) print(f'EPOCH: [{epoch}/{NUM_EPOCHS}]') print(f'TRAIN LOSS: {train_loss:.3f}, TRAIN ACC: {train_acc:.3f}') print(f'TEST LOSS: {test_loss:.3f}, TEST ACC: {test_acc:.3f}') # このように重みを保存する parameters = model.state_dict() torch.save(parameters, f'../weights/{epoch}.pth')
def main(): """ This code is written for the pre-training of the AlexNet to implement the SA-Siam object tarcker. SA-Siam has the two subnetwork, S-Net and A-Net, and this pre-trained AlexNet is used for the feature extractor of S-Net. I slightly changed the code from the pytorch examples (https://github.com/pytorch/examples/tree/master/imagenet) """ global args, best_prec1 args = parser.parse_args() args.distributed = args.world_size > 1 if args.distributed: dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url, world_size=args.world_size) # create model model = AlexNet() model = torch.nn.parallel.DataParallel(model).cuda() # model = torch.nn.parallel.DistributedDataParallel(model) # define loss function (criterion) and optimizer criterion = nn.CrossEntropyLoss().cuda() optimizer = torch.optim.SGD(model.parameters(), args.lr, momentum=args.momentum, weight_decay=args.weight_decay) # optionally resume from a checkpoint if args.resume: if os.path.isfile(args.resume): print("=> loading checkpoint '{}'".format(args.resume)) checkpoint = torch.load(args.resume) args.start_epoch = checkpoint['epoch'] best_prec1 = checkpoint['best_prec1'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) print("=> loaded checkpoint '{}' (epoch {})".format( args.resume, checkpoint['epoch'])) else: print("=> no checkpoint found at '{}'".format(args.resume)) cudnn.benchmark = True # Data loading code traindir = os.path.join(args.data, 'train') valdir = os.path.join(args.data, 'val') normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) train_dataset = datasets.ImageFolder( traindir, transforms.Compose([ transforms.RandomResizedCrop(255), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize, ])) if args.distributed: train_sampler = torch.utils.data.distributed.DistributedSampler( train_dataset) else: train_sampler = None train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=(train_sampler is None), num_workers=args.workers, pin_memory=True, sampler=train_sampler) val_loader = torch.utils.data.DataLoader( datasets.ImageFolder( valdir, transforms.Compose([ # transforms.Resize((255, 255)), transforms.RandomResizedCrop(255), # transforms.CenterCrop(255), transforms.ToTensor(), normalize, ])), batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True) if args.evaluate: validate(val_loader, model, criterion) return for epoch in range(args.start_epoch, args.epochs): if args.distributed: train_sampler.set_epoch(epoch) adjust_learning_rate(optimizer, epoch) # train for one epoch train(train_loader, model, criterion, optimizer, epoch) # evaluate on validation set prec1 = validate(val_loader, model, criterion) # remember best prec@1 and save checkpoint is_best = prec1 > best_prec1 best_prec1 = max(prec1, best_prec1) save_checkpoint( { 'epoch': epoch + 1, 'state_dict': model.state_dict(), 'best_prec1': best_prec1, 'optimizer': optimizer.state_dict(), }, is_best)