sneedium/captcha_pixelplanet

1import torch2import torch.nn as nn3from fastai.vision import *4 5from .model_vision import BaseIterVision6from .model_language import BCNLanguage7from .model_alignment import BaseAlignment8 9class IterNet(nn.Module):10    def __init__(self, config):11        super().__init__()12        self.iter_size = ifnone(config.model_iter_size, 1)13        self.max_length = config.dataset_max_length + 1  # additional stop token14        self.vision = BaseIterVision(config)15        self.language = BCNLanguage(config)16        self.alignment = BaseAlignment(config)17        self.deep_supervision = ifnone(config.model_deep_supervision, True)18 19    def forward(self, images, *args):20        list_v_res = self.vision(images)21        if not isinstance(list_v_res, (list, tuple)):22            list_v_res = [list_v_res]23        all_l_res, all_a_res = [], []24 25        for v_res in list_v_res:26            a_res = v_res27            for _ in range(self.iter_size):28                tokens = torch.softmax(a_res['logits'], dim=-1)29                lengths = a_res['pt_lengths']30                lengths.clamp_(2, self.max_length)  # TODO:move to langauge model31                l_res = self.language(tokens, lengths)32                all_l_res.append(l_res)33                a_res = self.alignment(l_res['feature'], v_res['feature'])34                all_a_res.append(a_res)35        if self.training and self.deep_supervision:36            return all_a_res, all_l_res, list_v_res37        else:38            return a_res, all_l_res[-1], list_v_res[-1]39