Pytorch上在最好的模型上继续训练模型（model 和 optimizer模型如何使用）


前情：在做实验的时候一开始用的是网上的代码，最后保存两个模型:model.pth,optimizer.pth，但是最终的实验结果和预期相差太大，准备在最好的模型基础上继续训练。

model.pth和optimizer.pth的区别

这里不讲这两个的具体作用，主要从用法开始讲：
1.如果你只是在已得到的最佳结果进行测试的话，只需要 model.pth 即可。
2.如果是想在已有的最佳结果上面继续训练，继续提升结果，那就需要到 model.pth 和optimizer.pth。

如何将已得到的模型参数（model.pth）和优化器(optimizer.pth)参数代入到工程代码中继续训练

需要修改以下几个地方：
1.在创建model之后加载model.pth (具体添加方法如以代码中的备注1)
2.在创建optimizer之后加载optimizer.pth(具体添加方法如代码中备注2)
3.给出model.pth和optimizer.pth的地址，也就是代码中的cfg.MODEL.CONTINUE.MODEL，cfg.MODEL.CONTINUE.OPTIMIZER，这个就是模型具体地址的位置。
4.非常重要，要将model从cpu转换到cuda，否则会报错

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
# encoding: utf-8

def train(cfg):
    model = build_fcn_model(cfg)
    #4.非常重要，要将model从cpu转换到cuda，否则会报错
    model=model.cuda()
    # 1.在model下面添加model.pth ,其中cfg.MODEL.CONTINUE.MODEL为模型保存的位置
    model.load_state_dict(torch.load(cfg.MODEL.CONTINUE.MODEL))
    optimizer = make_optimizer(cfg, model)
     # 2.在optimizer 下面添加optimizer .pth ,其中cfg.MODEL.CONTINUE.optimizer 为模型保存的位置
    optimizer.load_state_dict(torch.load(cfg.MODEL.CONTINUE.OPTIMIZER))
    arguments = {}

    data_loader = make_data_loader(cfg, is_train=True)
    val_loader = make_data_loader(cfg, is_train=False)

    do_train(
        cfg,
        model,
        data_loader,
        val_loader,
        optimizer,
        cross_entropy2d,
    )


def main():
    parser = argparse.ArgumentParser(description="PyTorch Training")
    parser.add_argument(
        "--config_file", default="", help="path to config file", type=str
    )
    parser.add_argument("opts", help="Modify config options using the command-line", default=None,
                        nargs=argparse.REMAINDER)

    args = parser.parse_args()

    num_gpus = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1

    if args.config_file != "":
        cfg.merge_from_file(args.config_file)
    cfg.merge_from_list(args.opts)
    cfg.freeze()

    output_dir = cfg.OUTPUT_DIR
    if output_dir and not os.path.exists(output_dir):
        mkdir(output_dir)

    logger = setup_logger("Model", output_dir, 0)
    logger.info("Using {} GPUS".format(num_gpus))
    logger.info(args)

    if args.config_file != "":
        logger.info("Loaded configuration file {}".format(args.config_file))
        with open(args.config_file, 'r') as cf:
            config_str = "\n" + cf.read()
            logger.info(config_str)
    logger.info("Running with config:\n{}".format(cfg))

    train(cfg)


if __name__ == '__main__':
    main()