VisualReasoning_MMnet/load_model.py at master · rs9000/VisualReasoning_MMnet · GitHub

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
from models.SAN import SAN as SAN
from models.SAN_wbw import SAN as SAN_wbw
from models.PG import PG
from models.PG_memory import PG as PG_memory
from models.PG_endtoend import PG as PG_endtoend


def load_model(args, vocab):

    modelname = args.model
    model = None

    if modelname == "SAN":
        model = SAN(vocab=vocab,
                    stem_dim=args.stem_dim,
                    question_size=args.question_size,
                    n_answers=args.answer_size,
                    batch_size=args.batch_size,
                    n_channel=args.n_channel
                    ).cuda()

    elif modelname == "SAN_wbw":
        model = SAN_wbw(vocab=vocab,
                        stem_dim=args.stem_dim,
                        question_size=args.question_size,
                        n_answers=args.answer_size,
                        batch_size=args.batch_size,
                        n_channel=args.n_channel
                        ).cuda()

    elif modelname == "PG":
        model = PG(vocab=vocab,
                   stem_dim=args.stem_dim,
                   question_size=args.question_size,
                   n_answers=args.answer_size,
                   batch_size=args.batch_size,
                   n_channel=args.n_channel
                   ).cuda()

    elif modelname == "PG_memory":
        model = PG_memory(vocab=vocab,
                          stem_dim=args.stem_dim,
                          question_size=args.question_size,
                          n_answers=args.answer_size,
                          batch_size=args.batch_size,
                          n_channel=args.n_channel
                          ).cuda()

    elif modelname == "PG_endtoend":
        model = PG_endtoend(vocab=vocab,
                            stem_dim=args.stem_dim,
                            question_size=args.question_size,
                            n_answers=args.answer_size,
                            batch_size=args.batch_size,
                            n_channel=args.n_channel,
                            decoder_mode=args.decoder_mode,
                            use_curriculum=args.use_curriculum
                            ).cuda()

    return model