自定义数据加载器

from torch.utils.data import Dataset
from torchvision import transforms
import torch
import os
from PIL import Image


class mydatas(Dataset):
    def __init__(self, path,image_size):
        self.file_path = path
        self.image_size = image_size
        self.images_path = []
        image_names_list = sorted(os.listdir(self.file_path))
        self.images_path = [os.path.join(self.file_path, name) for name in image_names_list]

    def __getitem__(self, idx):
        target = [0]
        transform = transforms.Compose([
            transforms.Resize([self.image_size,self.image_size]),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
        ])
        image_path = self.images_path[idx]
        image = Image.open(image_path).convert('RGB')
        image = transform(image)

        return image,target

    def __len__(self):
        return len(self.images_path)


if __name__=='__main__':
    dataset = mydatas('./data/FFA', image_size=128)

    data_loader = torch.utils.data.DataLoader(dataset,
                                              batch_size=1,
                                              shuffle=False,
                                              num_workers=4,
                                              pin_memory=True,
                                              drop_last=True)
    for iteration, (x,y) in enumerate(data_loader):
        print(x.shape)
        exit()

 保存权重到词典当中

 def save(self, dir, step):
        params = {}
        params['gen2B'] = self.gen2B.state_dict()
        params['gen2A'] = self.gen2A.state_dict()
        params['disA'] = self.disA.state_dict()
        params['disB'] = self.disB.state_dict()
        params['D_optimizer'] = self.D_optim.state_dict()
        params['G_optimizer'] = self.G_optim.state_dict()
        params['start_iter'] = step
        torch.save(params, os.path.join(dir, self.dataset + '_params_%07d.pt' % step))

tensor张量转为图片保存在指定文件夹

from torchvision.utils import save_image
def save_tensor_images(image_tensor, path, epoch):
    os.makedirs(path, exist_ok=True)
    for i in range(image_tensor.size(0)):
        save_image(image_tensor[i], os.path.join(path, f'image_{epoch}_{i}.jpg'))