def forward_pytorch(weightfile, image): #net=resnet.resnet18() #net = resnet.resnet18() net=LeNet(1,2) checkpoint = torch.load(weightfile) net.load_state_dict(checkpoint['weight']) net.double() # to double if args.cuda: net.cuda() print(net) net.eval() image = torch.from_numpy(image.astype(np.float64)) # to double if args.cuda: image = Variable(image.cuda()) else: image = Variable(image) t0 = time.time() blobs = net.forward(image) print(blobs.data.numpy().flatten()) t1 = time.time() return t1-t0, blobs, net, torch.from_numpy(blobs.data.numpy())
import torchvision as tv import torch.nn as nn import torch.optim as optim from torch.autograd import Variable as V from torchvision import transforms from lenet import LeNet from sobolev import SobolevLoss USE_SOBOLEV = False student = LeNet() teacher = LeNet() teacher.load_state_dict(th.load('teacher.pth')) student = student.cuda() teacher = teacher.cuda() transform_train = transforms.Compose([ transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)), ]) transform_test = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)), ]) trainset = tv.datasets.CIFAR10(root='./data',
def main(): global args, rank, world_size, best_prec1, dataset_len if args.dist == 1: rank, world_size = dist_init() else: rank = 0 world_size = 1 model = LeNet() model.cuda() param_copy = [ param.clone().type(torch.cuda.FloatTensor).detach() for param in model.parameters() ] for param in param_copy: param.requires_grad = True if args.dist == 1: model = DistModule(model) # define loss function (criterion) and optimizer criterion = nn.CrossEntropyLoss() optimizer = torch.optim.SGD(param_copy, args.base_lr, momentum=args.momentum, weight_decay=args.weight_decay) # optionally resume from a checkpoint last_iter = -1 # Data loading code train_dataset = datasets.MNIST(root='./data', train=True, transform=transforms.ToTensor(), download=False) val_dataset = datasets.MNIST(root='./data', train=False, transform=transforms.ToTensor(), download=False) dataset_len = len(train_dataset) args.max_iter = math.ceil( (dataset_len * args.epoch) / (world_size * args.batch_size)) if args.dist == 1: train_sampler = DistributedGivenIterationSampler(train_dataset, args.max_iter, args.batch_size, last_iter=last_iter) val_sampler = DistributedSampler(val_dataset, round_up=False) else: train_sampler = DistributedGivenIterationSampler(train_dataset, args.max_iter, args.batch_size, world_size=1, rank=0, last_iter=last_iter) val_sampler = None # pin_memory if true, will copy the tensor to cuda pinned memory train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True, sampler=train_sampler) val_loader = DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True, sampler=val_sampler) train(train_loader, val_loader, model, criterion, optimizer, param_copy)
#!/usr/bin/env python3 import torch as th import torchvision as tv import torch.nn as nn import torch.optim as optim from torch.autograd import Variable as V from torchvision import transforms from lenet import LeNet teacher = LeNet() teacher = teacher.cuda() transform_train = transforms.Compose([ transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)), ]) transform_test = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)), ]) trainset = tv.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_train)
w, h = shape trainloader = torch.utils.data.DataLoader(trainset, batch_size=args.batch_size, shuffle=True, pin_memory=cuda) testloader = torch.utils.data.DataLoader(testset, batch_size=args.batch_size, shuffle=False, pin_memory=cuda) model = LeNet(k=len(classes), args=args) if cuda: model.cuda() for epoch in range(1, args.epochs + 1): model.train() for i, (x_batch, y_batch) in enumerate(trainloader): if cuda: x_batch, y_batch = x_batch.cuda(), y_batch.cuda() model.optimizer.zero_grad() # forward + backward + optimize output = model(x_batch) loss = model.criterion(output, y_batch) loss.backward() model.optimizer.step()
filename = os.path.join(root, 'soft_targets.pickle') f = open(filename, 'wb') pickle.dump(self.train_labels, f) f.close() def get_soft_label(self, img): img = Image.fromarray(img.numpy(), mode='L') if self.transform is not None: img = self.transform(img) return self.init_target_transform(img) leNetModel = LeNet(args) if args.cuda: leNetModel.cuda() soft_labels = [] if importMode: leNetModel.load_state_dict( torch.load(os.path.join('./result', 'lenet5_best.json'))) else: for epoch in range(1, args.epochs + 1): leNetModel.train_(train_loader, epoch) def get_soft_label(img): return leNetModel.forward_with_temperature(img.view( 1, *(img.size()))).view(-1)