Pytorch-讀數據 or 模型

val_tf = pretrainedmodels.utils.TransformImage(
    model,
    scale=scale,
    preserve_aspect_ratio=args.preserve_aspect_ratio
)

val_loader = torch.utils.data.DataLoader(
    datasets.ImageFolder(valdir, val_tf),
    batch_size=args.batch_size, shuffle=False,
    num_workers=args.workers, pin_memory=True)

# 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)

model = torch.nn.DataParallel(model).cuda()

https://github.com/Cadene/pretrained-models.pytorch

https://www.e-learn.cn/content/qita/850153

 

 

# 保存模型

torch.save(model.state_dict(), alexnet.pth)

 

 

 

# 定義模型

model = AlexNet()

# 根據模型讀取參數

model.load_state_dict(state_dict)

# 權重信息

restore_param = {v for k, v in model.state_dict().items() }

# 網絡信息

model.parameters