Python torchvision 模块，transforms() 实例源码

def test():
    import torchvision

    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.485,0.456,0.406), (0.229,0.224,0.225))
    ])
    dataset = ListDataset(root='/mnt/hgfs/D/download/PASCAL_VOC/voc_all_images',
                          list_file='./data/voc12_train.txt', train=True, transform=transform, input_size=600)
    dataloader = torch.utils.data.DataLoader(dataset, batch_size=8, shuffle=False, num_workers=1, collate_fn=dataset.collate_fn)

    for images, loc_targets, cls_targets in dataloader:
        print(images.size())
        print(loc_targets.size())
        print(cls_targets.size())
        grid = torchvision.utils.make_grid(images, 1)
        torchvision.utils.save_image(grid, 'a.jpg')
        break

# test()

def load_data(resize):

    data_transforms = {
        'train': transforms.Compose([
            transforms.RandomSizedCrop(max(resize)),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ]),
        'val': transforms.Compose([
            #Higher scale-up for inception
            transforms.Scale(int(max(resize)/224*256)),
            transforms.CenterCrop(max(resize)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ]),
    }

    data_dir = 'PlantVillage'
    dsets = {x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x])
             for x in ['train', 'val']}
    dset_loaders = {x: torch.utils.data.DataLoader(dsets[x], batch_size=batch_size,
                                                   shuffle=True)
                    for x in ['train', 'val']}
    dset_sizes = {x: len(dsets[x]) for x in ['train', 'val']}
    dset_classes = dsets['train'].classes

    return dset_loaders['train'], dset_loaders['val']

def get_cifar10_loaders(root_directory, train_batch_size=128, test_batch_size=100,
                        download=False):
    # Data preparation for CIFAR10. Borrowed from
    # https://github.com/kuangliu/pytorch-cifar/blob/master/main.py
    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 = torchvision.datasets.CIFAR10(root=os.path.join(root_directory, 'data'),
                                            train=True, download=download,
                                            transform=transform_train)
    trainloader = torch.utils.data.DataLoader(trainset, batch_size=train_batch_size,
                                              shuffle=True, num_workers=2)

    testset = torchvision.datasets.CIFAR10(root=os.path.join(root_directory, 'data'),
                                           train=False, download=download,
                                           transform=transform_test)
    testloader = torch.utils.data.DataLoader(testset, batch_size=test_batch_size,
                                             shuffle=False, num_workers=2)
    return trainloader, testloader

def main():

    dir_data = './data'
    ext_img = 'png'
    #n_epoch = 100
    n_epoch = 50
    #n_img_per_batch = 40
    n_img_per_batch = 60
    #n_img_per_batch = 1
    n_worker = 4
    interval_train_loss = int(round(20000 / n_img_per_batch)) * n_img_per_batch
    is_gpu = torch.cuda.device_count() > 0

    transform = transforms.Compose(
        [transforms.ToTensor(),
         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

    di_set_transform = {'train' : transform, 'test' : transform}

    #fig = plt.figure(num=None, figsize=(1, 2), dpi=500)
    fig = plt.figure(num=None, figsize=(12, 18), dpi=100)
    plt.ion()
    ax_loss_train = fig.add_subplot(2, 1, 1)
    ax_loss_train.set_title('Avg. train loss per image vs. # train input images')
    ax_loss_train.xaxis.set_major_locator(MaxNLocator(integer=True))
    ax_loss_val = fig.add_subplot(2, 1, 2)
    ax_loss_val.set_title('Avg. val. loss per image vs. # train input images')
    ax_loss_val.xaxis.set_major_locator(MaxNLocator(integer=True))
    trainloader, testloader, net, criterion, optimizer, scheduler, li_class = \
        initialize(
            is_gpu, dir_data, di_set_transform, ext_img, n_img_per_batch, n_worker)
    #print('[%s] lap of initializing : %d sec' % (lap_sec))
    kolor = np.random.rand(3)
    #if 2 == i_m:
    #    a = 0
    train(is_gpu, trainloader, testloader, net, criterion, optimizer, scheduler, #li_class,
          n_epoch, ax_loss_train, ax_loss_val,
          kolor, interval_train_loss)
    print('Finished all.')
    plt.pause(1000)
    return

def __init__(self, opt):
        transform_list = []

        if (opt.crop_height > 0) and (opt.crop_width > 0):
            transform_list.append(transforms.CenterCrop(opt.crop_height, crop_width))
        elif opt.crop_size > 0:
            transform_list.append(transforms.CenterCrop(opt.crop_size))

        transform_list.append(transforms.Scale(opt.image_size))
        transform_list.append(transforms.CenterCrop(opt.image_size))

        transform_list.append(transforms.ToTensor())

        if opt.dataset == 'cifar10':
            dataset1 = datasets.CIFAR10(root = opt.dataroot, download = True,
                transform = transforms.Compose(transform_list))
            dataset2 = datasets.CIFAR10(root = opt.dataroot, train = False,
                transform = transforms.Compose(transform_list))
            def get_data(k):
                if k < len(dataset1):
                    return dataset1[k][0]
                else:
                    return dataset2[k - len(dataset1)][0]
        else:
            if opt.dataset in ['imagenet', 'folder', 'lfw']:
                dataset = datasets.ImageFolder(root = opt.dataroot,
                    transform = transforms.Compose(transform_list))
            elif opt.dataset == 'lsun':
                dataset = datasets.LSUN(db_path = opt.dataroot, classes = [opt.lsun_class + '_train'],
                    transform = transforms.Compose(transform_list))
            def get_data(k):
                return dataset[k][0]

        data_index = torch.load(os.path.join(opt.dataroot, 'data_index.pt'))
        train_index = data_index['train']

        self.opt = opt
        self.get_data = get_data
        self.train_index = data_index['train']
        self.counter = 0

def __init__(self, root, list_file, train, transform, input_size):
        '''
        Args:
          root: (str) ditectory to images.
          list_file: (str) path to index file.
          train: (boolean) train or test.
          transform: ([transforms]) image transforms.
          input_size: (int) model input size.
        '''
        self.root = root
        self.train = train
        self.transform = transform
        self.input_size = input_size

        self.fnames = []
        self.boxes = []
        self.labels = []

        self.encoder = DataEncoder()

        with open(list_file) as f:
            lines = f.readlines()
            self.num_samples = len(lines)

        for line in lines:
            splited = line.strip().split()
            self.fnames.append(splited[0])
            num_boxes = (len(splited) - 1) // 5
            box = []
            label = []
            for i in range(num_boxes):
                xmin = splited[1+5*i]
                ymin = splited[2+5*i]
                xmax = splited[3+5*i]
                ymax = splited[4+5*i]
                c = splited[5+5*i]
                box.append([float(xmin),float(ymin),float(xmax),float(ymax)])
                label.append(int(c))
            self.boxes.append(torch.Tensor(box))
            self.labels.append(torch.LongTensor(label))

def train():
    input = Variable(torch.FloatTensor(batch_size, num_history, num_components, image_size, image_size)).cuda()
    label = Variable(torch.FloatTensor(batch_size, num_components, image_size, image_size)).cuda()

    num_epochs = 25
    save_every_iteration = 100
    out_dir = "test"
    for epoch in range(num_epochs):
        for i, data in enumerate(loader):
            predictor.zero_grad()
            x, y = data

            transform = transforms.Compose([
                transforms.Normalize((0.0, 0.0, 0.0), (1.0, 1.0, 1.0))
                ])
            x = transform(x)
            y = transform(y)

            input.data.resize_(x.size()).copy_(x)
            label.data.resize_(y.size()).copy_(y)

            output = predictor(input)
            output_loss = loss(output, label)
            output_loss.backward()
            predictor.optimizer.step()

            print("[" + str(epoch) + "/ " + str(i) + "] Loss: " + str(output_loss.data[0]))

            if i % save_every_iteration == 0:
                img_outputs = predict_test_sequence()
                out_file = '%s/epoch_%03d.png' % (out_dir, epoch)
                print("saving to: " + out_file)
                vutils.save_image(torch.FloatTensor(initial_outputs), out_dir + "/initial.png")
                vutils.save_image(img_outputs.data, out_file)
                vutils.save_image(img_outputs.data, out_dir + "/latest.png")

        print("Saving model...")
        model_out_dir = "."
        torch.save(predictor.state_dict(), '%s/model_epoch_%d.pth' % (model_out_dir, epoch))

def init_cnn(self, opt):
        """Lazy initialization of preprocessor model in case we don't need any
        image preprocessing.
        """
        try:
            import torch
            self.use_cuda = (not opt.get('no_cuda', False)
                             and torch.cuda.is_available())
            self.torch = torch
        except ModuleNotFoundError:
            raise ModuleNotFoundError('Need to install Pytorch: go to pytorch.org')
        from torch.autograd import Variable
        import torchvision
        import torchvision.transforms as transforms
        import torch.nn as nn

        try:
            import h5py
            self.h5py = h5py
        except ModuleNotFoundError:
            raise ModuleNotFoundError('Need to install h5py')

        if 'image_mode' not in opt or 'image_size' not in opt:
            raise RuntimeError(
                'Need to add image arguments to opt. See '
                'parlai.core.params.ParlaiParser.add_image_args')
        self.image_mode = opt['image_mode']
        self.image_size = opt['image_size']
        self.crop_size = opt['image_cropsize']

        if self.use_cuda:
            print('[ Using CUDA ]')
            torch.cuda.set_device(opt.get('gpu', -1))

        cnn_type, layer_num = self.image_mode_switcher()

        # initialize the pretrained CNN using pytorch.
        CNN = getattr(torchvision.models, cnn_type)

        # cut off the additional layer.
        self.netCNN = nn.Sequential(
            *list(CNN(pretrained=True).children())[:layer_num])

        # initialize the transform function using torch vision.
        self.transform = transforms.Compose([
            transforms.Scale(self.image_size),
            transforms.CenterCrop(self.crop_size),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225])
        ])

        # container for single image
        self.xs = torch.zeros(1, 3, self.crop_size, self.crop_size)

        if self.use_cuda:
            self.netCNN.cuda()
            self.xs = self.xs.cuda()

        # make self.xs variable.
        self.xs = Variable(self.xs)

def main():

    #'''
    li_mode = ['TORCHVISION_MEMORY', 'TORCHVISION_IMAGEFOLDER',
               'CUSTOM_MEMORY', 'CUSTOM_FILE', 'CUSTOM_TENSORDATSET']
    #'''
    '''
    li_mode = ['CUSTOM_TENSORDATSET', 'TORCHVISION_MEMORY',
               'TORCHVISION_IMAGEFOLDER', 'CUSTOM_MEMORY', 'CUSTOM_FILE']
    '''
    dir_data = './data'
    ext_img = 'png'
    #n_epoch = 100
    n_epoch = 50
    #n_img_per_batch = 40
    n_img_per_batch = 60
    n_worker = 4
    interval_train_loss = int(round(20000 / n_img_per_batch)) * n_img_per_batch
    is_gpu = torch.cuda.device_count() > 0

    transform = transforms.Compose(
        [transforms.ToTensor(),
         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

    di_set_transform = {'train' : transform, 'test' : transform}

    #fig = plt.figure(num=None, figsize=(1, 2), dpi=500)
    fig = plt.figure(num=None, figsize=(12, 18), dpi=100)
    plt.ion()
    ax_time = fig.add_subplot(3, 1, 1)
    ax_time.set_title(
        'Elapsed time (sec.) of validation on 10k images vs. epoch. Note that value for epoch 0 is the elapsed time of init.')
    ax_time.xaxis.set_major_locator(MaxNLocator(integer=True))
    ax_loss_train = fig.add_subplot(3, 1, 2)
    ax_loss_train.set_title('Avg. train loss per image vs. # train input images')
    ax_loss_train.xaxis.set_major_locator(MaxNLocator(integer=True))
    ax_loss_val = fig.add_subplot(3, 1, 3)
    ax_loss_val.set_title('Avg. val. loss per image vs. # train input images')
    ax_loss_val.xaxis.set_major_locator(MaxNLocator(integer=True))
    for i_m, mode in enumerate(li_mode):
        start = time()
        trainloader, testloader, net, criterion, optimizer, scheduler, li_class = \
            initialize(
                mode, is_gpu, dir_data, di_set_transform, ext_img, n_img_per_batch, n_worker)
        lap_init = time() - start
        #print('[%s] lap of initializing : %d sec' % (lap_sec))
        kolor = np.random.rand(3)
        #if 2 == i_m:
        #    a = 0
        train(is_gpu, trainloader, testloader, net, criterion, optimizer, scheduler, #li_class,
              n_epoch, lap_init, ax_time, ax_loss_train, ax_loss_val,
              mode, kolor, n_img_per_batch, interval_train_loss)
    print('Finished all.')
    plt.pause(1000)
    return

def _generate_images(self, nb_batches, g_fp, r_idx, opt, show_info, queue):
        import torch
        import torch.nn as nn
        import torch.optim as optim
        import torchvision
        import torchvision.datasets as datasets
        import torchvision.transforms as transforms
        from torch.autograd import Variable

        #np.random.seed(42)
        #random.seed(42)
        #torch.manual_seed(42)

        gen = GeneratorLearnedInputSpace(opt.width, opt.height, opt.nfeature, opt.nlayer, opt.code_size, opt.norm, n_lis_layers=opt.r_iterations, upscaling=opt.g_upscaling)
        if show_info:
            print("G:", gen)
        gen.cuda()
        prefix = "last"
        gen.load_state_dict(torch.load(g_fp))
        gen.train()

        print("Generating images for checkpoint G'%s'..." % (g_fp,))
        #imgs_by_riter = [[] for _ in range(1+opt.r_iterations)]
        images_all = []
        for i in range(nb_batches):
            code = Variable(torch.randn(opt.batch_size, opt.code_size).cuda(), volatile=True)

            #for r_idx in range(1+opt.r_iterations):
            images, _ = gen(code, n_execute_lis_layers=r_idx)
            images_np = (images.data.cpu().numpy() * 255).astype(np.uint8).transpose((0, 2, 3, 1))

            #from scipy import misc
            #print(np.average(images[0]), np.min(images[0]), np.max(images[0]))
            #print(np.average(images_fixed[0]), np.min(images_fixed[0]), np.max(images_fixed[0]))
            #misc.imshow(list(images_np)[0])
            #misc.imshow(list(images_fixed)[0])

            #imgs_by_riter[r_idx].extend(list(images_np))
            images_all.extend(images_np)

        result_str = pickle.dumps({
            "g_fp": g_fp,
            "images": images_all
        }, protocol=-1)
        queue.put(result_str)