File size: 5,133 Bytes
a80d6bb
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c74a070
 
 
 
 
 
a80d6bb
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c74a070
 
a80d6bb
c74a070
a80d6bb
c74a070
a80d6bb
 
c74a070
 
a80d6bb
 
 
c74a070
a80d6bb
c74a070
a80d6bb
 
 
c74a070
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a80d6bb
 
c74a070
 
 
 
 
 
a80d6bb
c74a070
 
 
a80d6bb
 
 
c74a070
 
 
 
 
 
a80d6bb
c74a070
a80d6bb
 
 
 
 
c74a070
a80d6bb
c74a070
a80d6bb
 
 
 
c74a070
a80d6bb
 
 
 
 
c74a070
 
 
a80d6bb
c74a070
a80d6bb
c74a070
 
a80d6bb
c74a070
 
 
 
 
a80d6bb
 
c74a070
 
a80d6bb
 
 
 
 
 
 
c74a070
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
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
# Copyright 2019-present NAVER Corp.
# CC BY-NC-SA 3.0
# Available only for non-commercial use

import os, pdb
import torch
import torch.optim as optim

from tools import common, trainer
from tools.dataloader import *
from nets.patchnet import *
from nets.losses import *

default_net = "Quad_L2Net_ConfCFS()"

toy_db_debug = """SyntheticPairDataset(
    ImgFolder('imgs'), 
            'RandomScale(256,1024,can_upscale=True)', 
            'RandomTilting(0.5), PixelNoise(25)')"""

db_web_images = """SyntheticPairDataset(
    web_images, 
        'RandomScale(256,1024,can_upscale=True)',
        'RandomTilting(0.5), PixelNoise(25)')"""

db_aachen_images = """SyntheticPairDataset(
    aachen_db_images, 
        'RandomScale(256,1024,can_upscale=True)', 
        'RandomTilting(0.5), PixelNoise(25)')"""

db_aachen_style_transfer = """TransformedPairs(
    aachen_style_transfer_pairs,
            'RandomScale(256,1024,can_upscale=True), RandomTilting(0.5), PixelNoise(25)')"""

db_aachen_flow = "aachen_flow_pairs"

data_sources = dict(
    D=toy_db_debug,
    W=db_web_images,
    A=db_aachen_images,
    F=db_aachen_flow,
    S=db_aachen_style_transfer,
)

default_dataloader = """PairLoader(CatPairDataset(`data`),
    scale   = 'RandomScale(256,1024,can_upscale=True)',
    distort = 'ColorJitter(0.2,0.2,0.2,0.1)',
    crop    = 'RandomCrop(192)')"""

default_sampler = """NghSampler2(ngh=7, subq=-8, subd=1, pos_d=3, neg_d=5, border=16,
                            subd_neg=-8,maxpool_pos=True)"""

default_loss = """MultiLoss(
        1, ReliabilityLoss(`sampler`, base=0.5, nq=20),
        1, CosimLoss(N=`N`),
        1, PeakyLoss(N=`N`))"""


class MyTrainer(trainer.Trainer):
    """This class implements the network training.
    Below is the function I need to overload to explain how to do the backprop.
    """

    def forward_backward(self, inputs):
        output = self.net(imgs=[inputs.pop("img1"), inputs.pop("img2")])
        allvars = dict(inputs, **output)
        loss, details = self.loss_func(**allvars)
        if torch.is_grad_enabled():
            loss.backward()
        return loss, details


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser("Train R2D2")

    parser.add_argument("--data-loader", type=str, default=default_dataloader)
    parser.add_argument(
        "--train-data",
        type=str,
        default=list("WASF"),
        nargs="+",
        choices=set(data_sources.keys()),
    )
    parser.add_argument(
        "--net", type=str, default=default_net, help="network architecture"
    )

    parser.add_argument(
        "--pretrained", type=str, default="", help="pretrained model path"
    )
    parser.add_argument(
        "--save-path", type=str, required=True, help="model save_path path"
    )

    parser.add_argument("--loss", type=str, default=default_loss, help="loss function")
    parser.add_argument(
        "--sampler", type=str, default=default_sampler, help="AP sampler"
    )
    parser.add_argument(
        "--N", type=int, default=16, help="patch size for repeatability"
    )

    parser.add_argument(
        "--epochs", type=int, default=25, help="number of training epochs"
    )
    parser.add_argument("--batch-size", "--bs", type=int, default=8, help="batch size")
    parser.add_argument("--learning-rate", "--lr", type=str, default=1e-4)
    parser.add_argument("--weight-decay", "--wd", type=float, default=5e-4)

    parser.add_argument(
        "--threads", type=int, default=8, help="number of worker threads"
    )
    parser.add_argument("--gpu", type=int, nargs="+", default=[0], help="-1 for CPU")

    args = parser.parse_args()

    iscuda = common.torch_set_gpu(args.gpu)
    common.mkdir_for(args.save_path)

    # Create data loader
    from datasets import *

    db = [data_sources[key] for key in args.train_data]
    db = eval(args.data_loader.replace("`data`", ",".join(db)).replace("\n", ""))
    print("Training image database =", db)
    loader = threaded_loader(db, iscuda, args.threads, args.batch_size, shuffle=True)

    # create network
    print("\n>> Creating net = " + args.net)
    net = eval(args.net)
    print(f" ( Model size: {common.model_size(net)/1000:.0f}K parameters )")

    # initialization
    if args.pretrained:
        checkpoint = torch.load(args.pretrained, lambda a, b: a)
        net.load_pretrained(checkpoint["state_dict"])

    # create losses
    loss = args.loss.replace("`sampler`", args.sampler).replace("`N`", str(args.N))
    print("\n>> Creating loss = " + loss)
    loss = eval(loss.replace("\n", ""))

    # create optimizer
    optimizer = optim.Adam(
        [p for p in net.parameters() if p.requires_grad],
        lr=args.learning_rate,
        weight_decay=args.weight_decay,
    )

    train = MyTrainer(net, loader, loss, optimizer)
    if iscuda:
        train = train.cuda()

    # Training loop #
    for epoch in range(args.epochs):
        print(f"\n>> Starting epoch {epoch}...")
        train()

    print(f"\n>> Saving model to {args.save_path}")
    torch.save({"net": args.net, "state_dict": net.state_dict()}, args.save_path)