├── .idea
├── inspectionProfiles
│ └── Project_Default.xml
├── modules.xml
└── translating-images-into-maps.iml
├── LICENSE.md
├── README.md
├── images
├── ._image_to_bev_motivation.gif
└── image_to_bev_motivation.gif
├── src
├── __init__.py
├── data
│ ├── augmentations.py
│ ├── collate_funcs.py
│ ├── data_generation.py
│ ├── dataloader.py
│ ├── dataloader_new.py
│ ├── nuscenes
│ │ ├── __init__.py
│ │ ├── dataset.py
│ │ ├── splits.py
│ │ └── utils.py
│ └── utils.py
├── model
│ ├── __init__.py
│ ├── axial_attention.py
│ ├── backbone_utils.py
│ ├── bev_transform.py
│ ├── dla.py
│ ├── dla_up.py
│ ├── fpn.py
│ ├── general_modules.py
│ ├── intermediate_layer_getter.py
│ ├── loss.py
│ ├── mocha.py
│ ├── mocha2.py
│ ├── monotonic_attention.py
│ ├── monotonic_attention00.py
│ ├── monotonic_attention_fairseq.py
│ ├── network.py
│ ├── positional_encoding.py
│ ├── resnet.py
│ └── transformer.py
└── utils.py
└── train.py
/.idea/inspectionProfiles/Project_Default.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
43 |
44 |
45 |
46 |
47 |
48 |
49 |
78 |
79 |
80 |
83 |
84 |
85 |
87 |
91 |
92 |
--------------------------------------------------------------------------------
/.idea/modules.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
7 |
8 |
--------------------------------------------------------------------------------
/.idea/translating-images-into-maps.iml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
8 |
--------------------------------------------------------------------------------
/LICENSE.md:
--------------------------------------------------------------------------------
1 | # Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International
2 |
3 | Creative Commons Corporation (“Creative Commons”) is not a law firm and does not provide legal services or legal advice. Distribution of Creative Commons public licenses does not create a lawyer-client or other relationship. Creative Commons makes its licenses and related information available on an “as-is” basis. Creative Commons gives no warranties regarding its licenses, any material licensed under their terms and conditions, or any related information. Creative Commons disclaims all liability for damages resulting from their use to the fullest extent possible.
4 |
5 | ### Using Creative Commons Public Licenses
6 |
7 | Creative Commons public licenses provide a standard set of terms and conditions that creators and other rights holders may use to share original works of authorship and other material subject to copyright and certain other rights specified in the public license below. The following considerations are for informational purposes only, are not exhaustive, and do not form part of our licenses.
8 |
9 | * __Considerations for licensors:__ Our public licenses are intended for use by those authorized to give the public permission to use material in ways otherwise restricted by copyright and certain other rights. Our licenses are irrevocable. Licensors should read and understand the terms and conditions of the license they choose before applying it. Licensors should also secure all rights necessary before applying our licenses so that the public can reuse the material as expected. Licensors should clearly mark any material not subject to the license. This includes other CC-licensed material, or material used under an exception or limitation to copyright. [More considerations for licensors](http://wiki.creativecommons.org/Considerations_for_licensors_and_licensees#Considerations_for_licensors).
10 |
11 | * __Considerations for the public:__ By using one of our public licenses, a licensor grants the public permission to use the licensed material under specified terms and conditions. If the licensor’s permission is not necessary for any reason–for example, because of any applicable exception or limitation to copyright–then that use is not regulated by the license. Our licenses grant only permissions under copyright and certain other rights that a licensor has authority to grant. Use of the licensed material may still be restricted for other reasons, including because others have copyright or other rights in the material. A licensor may make special requests, such as asking that all changes be marked or described. Although not required by our licenses, you are encouraged to respect those requests where reasonable. [More considerations for the public](http://wiki.creativecommons.org/Considerations_for_licensors_and_licensees#Considerations_for_licensees).
12 |
13 | ## Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International Public License
14 |
15 | By exercising the Licensed Rights (defined below), You accept and agree to be bound by the terms and conditions of this Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International Public License ("Public License"). To the extent this Public License may be interpreted as a contract, You are granted the Licensed Rights in consideration of Your acceptance of these terms and conditions, and the Licensor grants You such rights in consideration of benefits the Licensor receives from making the Licensed Material available under these terms and conditions.
16 |
17 | ### Section 1 – Definitions.
18 |
19 | a. __Adapted Material__ means material subject to Copyright and Similar Rights that is derived from or based upon the Licensed Material and in which the Licensed Material is translated, altered, arranged, transformed, or otherwise modified in a manner requiring permission under the Copyright and Similar Rights held by the Licensor. For purposes of this Public License, where the Licensed Material is a musical work, performance, or sound recording, Adapted Material is always produced where the Licensed Material is synched in timed relation with a moving image.
20 |
21 | b. __Adapter's License__ means the license You apply to Your Copyright and Similar Rights in Your contributions to Adapted Material in accordance with the terms and conditions of this Public License.
22 |
23 | c. __BY-NC-SA Compatible License__ means a license listed at [creativecommons.org/compatiblelicenses](http://creativecommons.org/compatiblelicenses), approved by Creative Commons as essentially the equivalent of this Public License.
24 |
25 | d. __Copyright and Similar Rights__ means copyright and/or similar rights closely related to copyright including, without limitation, performance, broadcast, sound recording, and Sui Generis Database Rights, without regard to how the rights are labeled or categorized. For purposes of this Public License, the rights specified in Section 2(b)(1)-(2) are not Copyright and Similar Rights.
26 |
27 | e. __Effective Technological Measures__ means those measures that, in the absence of proper authority, may not be circumvented under laws fulfilling obligations under Article 11 of the WIPO Copyright Treaty adopted on December 20, 1996, and/or similar international agreements.
28 |
29 | f. __Exceptions and Limitations__ means fair use, fair dealing, and/or any other exception or limitation to Copyright and Similar Rights that applies to Your use of the Licensed Material.
30 |
31 | g. __License Elements__ means the license attributes listed in the name of a Creative Commons Public License. The License Elements of this Public License are Attribution, NonCommercial, and ShareAlike.
32 |
33 | h. __Licensed Material__ means the artistic or literary work, database, or other material to which the Licensor applied this Public License.
34 |
35 | i. __Licensed Rights__ means the rights granted to You subject to the terms and conditions of this Public License, which are limited to all Copyright and Similar Rights that apply to Your use of the Licensed Material and that the Licensor has authority to license.
36 |
37 | h. __Licensor__ means the individual(s) or entity(ies) granting rights under this Public License.
38 |
39 | i. __NonCommercial__ means not primarily intended for or directed towards commercial advantage or monetary compensation. For purposes of this Public License, the exchange of the Licensed Material for other material subject to Copyright and Similar Rights by digital file-sharing or similar means is NonCommercial provided there is no payment of monetary compensation in connection with the exchange.
40 |
41 | j. __Share__ means to provide material to the public by any means or process that requires permission under the Licensed Rights, such as reproduction, public display, public performance, distribution, dissemination, communication, or importation, and to make material available to the public including in ways that members of the public may access the material from a place and at a time individually chosen by them.
42 |
43 | k. __Sui Generis Database Rights__ means rights other than copyright resulting from Directive 96/9/EC of the European Parliament and of the Council of 11 March 1996 on the legal protection of databases, as amended and/or succeeded, as well as other essentially equivalent rights anywhere in the world.
44 |
45 | l. __You__ means the individual or entity exercising the Licensed Rights under this Public License. Your has a corresponding meaning.
46 |
47 | ### Section 2 – Scope.
48 |
49 | a. ___License grant.___
50 |
51 | 1. Subject to the terms and conditions of this Public License, the Licensor hereby grants You a worldwide, royalty-free, non-sublicensable, non-exclusive, irrevocable license to exercise the Licensed Rights in the Licensed Material to:
52 |
53 | A. reproduce and Share the Licensed Material, in whole or in part, for NonCommercial purposes only; and
54 |
55 | B. produce, reproduce, and Share Adapted Material for NonCommercial purposes only.
56 |
57 | 2. __Exceptions and Limitations.__ For the avoidance of doubt, where Exceptions and Limitations apply to Your use, this Public License does not apply, and You do not need to comply with its terms and conditions.
58 |
59 | 3. __Term.__ The term of this Public License is specified in Section 6(a).
60 |
61 | 4. __Media and formats; technical modifications allowed.__ The Licensor authorizes You to exercise the Licensed Rights in all media and formats whether now known or hereafter created, and to make technical modifications necessary to do so. The Licensor waives and/or agrees not to assert any right or authority to forbid You from making technical modifications necessary to exercise the Licensed Rights, including technical modifications necessary to circumvent Effective Technological Measures. For purposes of this Public License, simply making modifications authorized by this Section 2(a)(4) never produces Adapted Material.
62 |
63 | 5. __Downstream recipients.__
64 |
65 | A. __Offer from the Licensor – Licensed Material.__ Every recipient of the Licensed Material automatically receives an offer from the Licensor to exercise the Licensed Rights under the terms and conditions of this Public License.
66 |
67 | B. __Additional offer from the Licensor – Adapted Material.__ Every recipient of Adapted Material from You automatically receives an offer from the Licensor to exercise the Licensed Rights in the Adapted Material under the conditions of the Adapter’s License You apply.
68 |
69 | C. __No downstream restrictions.__ You may not offer or impose any additional or different terms or conditions on, or apply any Effective Technological Measures to, the Licensed Material if doing so restricts exercise of the Licensed Rights by any recipient of the Licensed Material.
70 |
71 | 6. __No endorsement.__ Nothing in this Public License constitutes or may be construed as permission to assert or imply that You are, or that Your use of the Licensed Material is, connected with, or sponsored, endorsed, or granted official status by, the Licensor or others designated to receive attribution as provided in Section 3(a)(1)(A)(i).
72 |
73 | b. ___Other rights.___
74 |
75 | 1. Moral rights, such as the right of integrity, are not licensed under this Public License, nor are publicity, privacy, and/or other similar personality rights; however, to the extent possible, the Licensor waives and/or agrees not to assert any such rights held by the Licensor to the limited extent necessary to allow You to exercise the Licensed Rights, but not otherwise.
76 |
77 | 2. Patent and trademark rights are not licensed under this Public License.
78 |
79 | 3. To the extent possible, the Licensor waives any right to collect royalties from You for the exercise of the Licensed Rights, whether directly or through a collecting society under any voluntary or waivable statutory or compulsory licensing scheme. In all other cases the Licensor expressly reserves any right to collect such royalties, including when the Licensed Material is used other than for NonCommercial purposes.
80 |
81 | ### Section 3 – License Conditions.
82 |
83 | Your exercise of the Licensed Rights is expressly made subject to the following conditions.
84 |
85 | a. ___Attribution.___
86 |
87 | 1. If You Share the Licensed Material (including in modified form), You must:
88 |
89 | A. retain the following if it is supplied by the Licensor with the Licensed Material:
90 |
91 | i. identification of the creator(s) of the Licensed Material and any others designated to receive attribution, in any reasonable manner requested by the Licensor (including by pseudonym if designated);
92 |
93 | ii. a copyright notice;
94 |
95 | iii. a notice that refers to this Public License;
96 |
97 | iv. a notice that refers to the disclaimer of warranties;
98 |
99 | v. a URI or hyperlink to the Licensed Material to the extent reasonably practicable;
100 |
101 | B. indicate if You modified the Licensed Material and retain an indication of any previous modifications; and
102 |
103 | C. indicate the Licensed Material is licensed under this Public License, and include the text of, or the URI or hyperlink to, this Public License.
104 |
105 | 2. You may satisfy the conditions in Section 3(a)(1) in any reasonable manner based on the medium, means, and context in which You Share the Licensed Material. For example, it may be reasonable to satisfy the conditions by providing a URI or hyperlink to a resource that includes the required information.
106 |
107 | 3. If requested by the Licensor, You must remove any of the information required by Section 3(a)(1)(A) to the extent reasonably practicable.
108 |
109 | b. ___ShareAlike.___
110 |
111 | In addition to the conditions in Section 3(a), if You Share Adapted Material You produce, the following conditions also apply.
112 |
113 | 1. The Adapter’s License You apply must be a Creative Commons license with the same License Elements, this version or later, or a BY-NC-SA Compatible License.
114 |
115 | 2. You must include the text of, or the URI or hyperlink to, the Adapter's License You apply. You may satisfy this condition in any reasonable manner based on the medium, means, and context in which You Share Adapted Material.
116 |
117 | 3. You may not offer or impose any additional or different terms or conditions on, or apply any Effective Technological Measures to, Adapted Material that restrict exercise of the rights granted under the Adapter's License You apply.
118 |
119 | ### Section 4 – Sui Generis Database Rights.
120 |
121 | Where the Licensed Rights include Sui Generis Database Rights that apply to Your use of the Licensed Material:
122 |
123 | a. for the avoidance of doubt, Section 2(a)(1) grants You the right to extract, reuse, reproduce, and Share all or a substantial portion of the contents of the database for NonCommercial purposes only;
124 |
125 | b. if You include all or a substantial portion of the database contents in a database in which You have Sui Generis Database Rights, then the database in which You have Sui Generis Database Rights (but not its individual contents) is Adapted Material, including for purposes of Section 3(b); and
126 |
127 | c. You must comply with the conditions in Section 3(a) if You Share all or a substantial portion of the contents of the database.
128 |
129 | For the avoidance of doubt, this Section 4 supplements and does not replace Your obligations under this Public License where the Licensed Rights include other Copyright and Similar Rights.
130 |
131 | ### Section 5 – Disclaimer of Warranties and Limitation of Liability.
132 |
133 | a. __Unless otherwise separately undertaken by the Licensor, to the extent possible, the Licensor offers the Licensed Material as-is and as-available, and makes no representations or warranties of any kind concerning the Licensed Material, whether express, implied, statutory, or other. This includes, without limitation, warranties of title, merchantability, fitness for a particular purpose, non-infringement, absence of latent or other defects, accuracy, or the presence or absence of errors, whether or not known or discoverable. Where disclaimers of warranties are not allowed in full or in part, this disclaimer may not apply to You.__
134 |
135 | b. __To the extent possible, in no event will the Licensor be liable to You on any legal theory (including, without limitation, negligence) or otherwise for any direct, special, indirect, incidental, consequential, punitive, exemplary, or other losses, costs, expenses, or damages arising out of this Public License or use of the Licensed Material, even if the Licensor has been advised of the possibility of such losses, costs, expenses, or damages. Where a limitation of liability is not allowed in full or in part, this limitation may not apply to You.__
136 |
137 | c. The disclaimer of warranties and limitation of liability provided above shall be interpreted in a manner that, to the extent possible, most closely approximates an absolute disclaimer and waiver of all liability.
138 |
139 | ### Section 6 – Term and Termination.
140 |
141 | a. This Public License applies for the term of the Copyright and Similar Rights licensed here. However, if You fail to comply with this Public License, then Your rights under this Public License terminate automatically.
142 |
143 | b. Where Your right to use the Licensed Material has terminated under Section 6(a), it reinstates:
144 |
145 | 1. automatically as of the date the violation is cured, provided it is cured within 30 days of Your discovery of the violation; or
146 |
147 | 2. upon express reinstatement by the Licensor.
148 |
149 | For the avoidance of doubt, this Section 6(b) does not affect any right the Licensor may have to seek remedies for Your violations of this Public License.
150 |
151 | c. For the avoidance of doubt, the Licensor may also offer the Licensed Material under separate terms or conditions or stop distributing the Licensed Material at any time; however, doing so will not terminate this Public License.
152 |
153 | d. Sections 1, 5, 6, 7, and 8 survive termination of this Public License.
154 |
155 | ### Section 7 – Other Terms and Conditions.
156 |
157 | a. The Licensor shall not be bound by any additional or different terms or conditions communicated by You unless expressly agreed.
158 |
159 | b. Any arrangements, understandings, or agreements regarding the Licensed Material not stated herein are separate from and independent of the terms and conditions of this Public License.
160 |
161 | ### Section 8 – Interpretation.
162 |
163 | a. For the avoidance of doubt, this Public License does not, and shall not be interpreted to, reduce, limit, restrict, or impose conditions on any use of the Licensed Material that could lawfully be made without permission under this Public License.
164 |
165 | b. To the extent possible, if any provision of this Public License is deemed unenforceable, it shall be automatically reformed to the minimum extent necessary to make it enforceable. If the provision cannot be reformed, it shall be severed from this Public License without affecting the enforceability of the remaining terms and conditions.
166 |
167 | c. No term or condition of this Public License will be waived and no failure to comply consented to unless expressly agreed to by the Licensor.
168 |
169 | d. Nothing in this Public License constitutes or may be interpreted as a limitation upon, or waiver of, any privileges and immunities that apply to the Licensor or You, including from the legal processes of any jurisdiction or authority.
170 |
171 | > Creative Commons is not a party to its public licenses. Notwithstanding, Creative Commons may elect to apply one of its public licenses to material it publishes and in those instances will be considered the “Licensor.” Except for the limited purpose of indicating that material is shared under a Creative Commons public license or as otherwise permitted by the Creative Commons policies published at [creativecommons.org/policies](http://creativecommons.org/policies), Creative Commons does not authorize the use of the trademark “Creative Commons” or any other trademark or logo of Creative Commons without its prior written consent including, without limitation, in connection with any unauthorized modifications to any of its public licenses or any other arrangements, understandings, or agreements concerning use of licensed material. For the avoidance of doubt, this paragraph does not form part of the public licenses.
172 | >
173 | > Creative Commons may be contacted at creativecommons.org
174 |
--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
1 | # Translating Images into Maps
2 | #### Avishkar Saha, Oscar Mendez Maldonado, Chris Russell and Richard Bowden
3 |
4 | This is the official code for the paper [Translating Images Into Maps](https://arxiv.org/abs/2110.00966) presented at ICRA 2022.
5 |
6 | ### Translating Images into Maps
7 |
8 |
9 |
`_.
37 | The feature maps are currently supposed to be in increasing depth
38 | order.
39 | The input to the model is expected to be an OrderedDict[Tensor], containing
40 | the feature maps on top of which the FPN will be added.
41 | Arguments:
42 | in_channels_list (list[int]): number of channels for each feature map that
43 | is passed to the module
44 | out_channels (int): number of channels of the FPN representation
45 | extra_blocks (ExtraFPNBlock or None): if provided, extra operations will
46 | be performed. It is expected to take the fpn features, the original
47 | features and the names of the original features as input, and returns
48 | a new list of feature maps and their corresponding names
49 | Examples::
50 | >>> m = torchvision.ops.FeaturePyramidNetwork([10, 20, 30], 5)
51 | >>> # get some dummy data
52 | >>> x = OrderedDict()
53 | >>> x['feat0'] = torch.rand(1, 10, 64, 64)
54 | >>> x['feat2'] = torch.rand(1, 20, 16, 16)
55 | >>> x['feat3'] = torch.rand(1, 30, 8, 8)
56 | >>> # compute the FPN on top of x
57 | >>> output = m(x)
58 | >>> print([(k, v.shape) for k, v in output.items()])
59 | >>> # returns
60 | >>> [('feat0', torch.Size([1, 5, 64, 64])),
61 | >>> ('feat2', torch.Size([1, 5, 16, 16])),
62 | >>> ('feat3', torch.Size([1, 5, 8, 8]))]
63 | """
64 |
65 | def __init__(
66 | self,
67 | in_channels_list: List[int],
68 | out_channels: int,
69 | extra_blocks: Optional[ExtraFPNBlock] = None,
70 | ):
71 | super(FeaturePyramidNetworkRoddick, self).__init__()
72 | self.inner_blocks = nn.ModuleList()
73 | self.layer_blocks = nn.ModuleList()
74 | for in_channels in in_channels_list:
75 | if in_channels == 0:
76 | raise ValueError("in_channels=0 is currently not supported")
77 | inner_block_module = nn.Conv2d(in_channels, out_channels, 1)
78 | layer_block_module = nn.Conv2d(out_channels, out_channels, 3, padding=1)
79 | self.inner_blocks.append(inner_block_module)
80 | self.layer_blocks.append(layer_block_module)
81 |
82 | # initialize parameters now to avoid modifying the initialization of top_blocks
83 | for m in self.children():
84 | if isinstance(m, nn.Conv2d):
85 | nn.init.kaiming_uniform_(m.weight, a=1)
86 | nn.init.constant_(m.bias, 0)
87 |
88 | if extra_blocks is not None:
89 | assert isinstance(extra_blocks, ExtraFPNBlock)
90 | self.extra_blocks = extra_blocks
91 |
92 | def get_result_from_inner_blocks(self, x: Tensor, idx: int) -> Tensor:
93 | """
94 | This is equivalent to self.inner_blocks[idx](x),
95 | but torchscript doesn't support this yet
96 | """
97 | num_blocks = 0
98 | for m in self.inner_blocks:
99 | num_blocks += 1
100 | if idx < 0:
101 | idx += num_blocks
102 | i = 0
103 | out = x
104 | for module in self.inner_blocks:
105 | if i == idx:
106 | out = module(x)
107 | i += 1
108 | return out
109 |
110 | def get_result_from_layer_blocks(self, x: Tensor, idx: int) -> Tensor:
111 | """
112 | This is equivalent to self.layer_blocks[idx](x),
113 | but torchscript doesn't support this yet
114 | """
115 | num_blocks = 0
116 | for m in self.layer_blocks:
117 | num_blocks += 1
118 | if idx < 0:
119 | idx += num_blocks
120 | i = 0
121 | out = x
122 | for module in self.layer_blocks:
123 | if i == idx:
124 | out = module(x)
125 | i += 1
126 | return out
127 |
128 | def forward(self, x: Dict[str, Tensor]) -> Dict[str, Tensor]:
129 | """
130 | Computes the FPN for a set of feature maps.
131 | Arguments:
132 | x (OrderedDict[Tensor]): feature maps for each feature level.
133 | Returns:
134 | results (OrderedDict[Tensor]): feature maps after FPN layers.
135 | They are ordered from highest resolution first.
136 | """
137 | # unpack OrderedDict into two lists for easier handling
138 | names = list(x.keys())
139 | x = list(x.values())
140 |
141 | last_inner = self.get_result_from_inner_blocks(x[-1], -1)
142 | results = []
143 | results.append(self.get_result_from_layer_blocks(last_inner, -1))
144 |
145 | for idx in range(len(x) - 2, -1, -1):
146 | inner_lateral = self.get_result_from_inner_blocks(x[idx], idx)
147 | feat_shape = inner_lateral.shape[-2:]
148 | inner_top_down = F.interpolate(last_inner, size=feat_shape, mode="nearest")
149 | last_inner = inner_lateral + inner_top_down
150 | results.insert(0, self.get_result_from_layer_blocks(last_inner, idx))
151 |
152 | if self.extra_blocks is not None:
153 | results, names = self.extra_blocks(results, x, names)
154 |
155 | # make it back an OrderedDict
156 | out = OrderedDict([(k, v) for k, v in zip(names, results)])
157 |
158 | return out
159 |
160 |
161 | class LLMaxPool(ExtraFPNBlock):
162 | """
163 | Applies a max_pool2d on top of the last feature map
164 | """
165 |
166 | def forward(
167 | self,
168 | x: List[Tensor],
169 | y: List[Tensor],
170 | names: List[str],
171 | ) -> Tuple[List[Tensor], List[str]]:
172 | names.append("pool")
173 | x.append(F.max_pool2d(x[-1], 1, 2, 0))
174 | return x, names
175 |
176 |
177 | class LastLevelP6P7(ExtraFPNBlock):
178 | """
179 | This module is used in RetinaNet to generate extra layers, P6 and P7.
180 | """
181 |
182 | def __init__(self, in_channels: int, out_channels: int):
183 | super(LastLevelP6P7, self).__init__()
184 | self.p6 = nn.Conv2d(in_channels, out_channels, 3, 2, 1)
185 | self.p7 = nn.Conv2d(out_channels, out_channels, 3, 2, 1)
186 | for module in [self.p6, self.p7]:
187 | nn.init.kaiming_uniform_(module.weight, a=1)
188 | nn.init.constant_(module.bias, 0)
189 | self.use_P5 = in_channels == out_channels
190 |
191 | def forward(
192 | self,
193 | p: List[Tensor],
194 | c: List[Tensor],
195 | names: List[str],
196 | ) -> Tuple[List[Tensor], List[str]]:
197 | p5, c5 = p[-1], c[-1]
198 | x = p5 if self.use_P5 else c5
199 | p6 = self.p6(x)
200 | p7 = self.p7(F.relu(p6))
201 | p.extend([p6, p7])
202 | names.extend(["p6", "p7"])
203 | return p, names
204 |
--------------------------------------------------------------------------------
/src/model/intermediate_layer_getter.py:
--------------------------------------------------------------------------------
1 | import functools
2 | from collections import OrderedDict
3 |
4 |
5 | # using wonder's beautiful simplification: https://stackoverflow.com/questions/31174295/getattr-and-setattr-on-nested-objects/31174427?noredirect=1#comment86638618_31174427
6 | def rgetattr(obj, attr, *args):
7 | def _getattr(obj, attr):
8 | return getattr(obj, attr, *args)
9 |
10 | return functools.reduce(_getattr, [obj] + attr.split("."))
11 |
12 |
13 | class IntermediateLayerGetter:
14 | def __init__(self, model, return_layers, keep_output=True):
15 | """Wraps a Pytorch module to get intermediate values
16 |
17 | Arguments:
18 | model {nn.module} -- The Pytorch module to call
19 | return_layers {dict} -- Dictionary with the selected submodules
20 | to return the output (format: {[current_module_name]: [desired_output_name]},
21 | current_module_name can be a nested submodule, e.g. submodule1.submodule2.submodule3)
22 |
23 | Keyword Arguments:
24 | keep_output {bool} -- If True model_output contains the final model's output
25 | in the other case model_output is None (default: {True})
26 | Returns:
27 | (mid_outputs {OrderedDict}, model_output {any}) -- mid_outputs keys are
28 | your desired_output_name (s) and their values are the returned tensors
29 | of those submodules (OrderedDict([(desired_output_name,tensor(...)), ...).
30 | See keep_output argument for model_output description.
31 | In case a submodule is called more than one time, all it's outputs are
32 | stored in a list.
33 | """
34 | self._model = model
35 | self.return_layers = return_layers
36 | self.keep_output = keep_output
37 |
38 | def __call__(self, *args, **kwargs):
39 | ret = OrderedDict()
40 | handles = []
41 | for name, new_name in self.return_layers.items():
42 | layer = rgetattr(self._model, name)
43 |
44 | def hook(module, input, output, new_name=new_name):
45 | if new_name in ret:
46 | if type(ret[new_name]) is list:
47 | ret[new_name].append(output)
48 | else:
49 | ret[new_name] = [ret[new_name], output]
50 | else:
51 | ret[new_name] = output
52 |
53 | try:
54 | h = layer.register_forward_hook(hook)
55 | except AttributeError as e:
56 | raise AttributeError(f"Module {name} not found")
57 | handles.append(h)
58 |
59 | if self.keep_output:
60 | output = self._model(*args, **kwargs)
61 | else:
62 | self._model(*args, **kwargs)
63 | output = None
64 |
65 | for h in handles:
66 | h.remove()
67 |
68 | return ret, output
69 |
--------------------------------------------------------------------------------
/src/model/loss.py:
--------------------------------------------------------------------------------
1 | import math
2 | import torch
3 | import torch.nn as nn
4 | import torch.nn.functional as F
5 | import numpy as np
6 | import matplotlib.pyplot as plt
7 |
8 | import src
9 | from torch.nn.utils.rnn import pad_sequence
10 |
11 |
12 | def class_obj_size():
13 | """
14 | NuScenes class object size on Roddick Split
15 | gt_dir/name: semantic_maps_new_200x200
16 | """
17 |
18 | # Object size per class relative to map resolution
19 | s200 = [0.263, 0.031, 0.055, 0.029, 0.009, 0.002, 0.010, 0.001, 0.008, 0.006, 0.001]
20 | s100 = [0.296, 0.037, 0.066, 0.035, 0.011, 0.003, 0.013, 0.002, 0.010, 0.008, 0.002]
21 | s50 = [0.331, 0.046, 0.083, 0.044, 0.015, 0.004, 0.019, 0.003, 0.014, 0.011, 0.003]
22 | s25 = [0.380, 0.065, 0.114, 0.061, 0.022, 0.007, 0.031, 0.006, 0.022, 0.017, 0.008]
23 | s13 = [0.465, 0.105, 0.169, 0.093, 0.037, 0.016, 0.058, 0.013, 0.039, 0.031, 0.021]
24 | s7 = [0.599, 0.196, 0.279, 0.148, 0.069, 0.037, 0.117, 0.033, 0.076, 0.065, 0.054]
25 |
26 | ms_obj_size = torch.stack(
27 | [
28 | torch.tensor(s100),
29 | torch.tensor(s50),
30 | torch.tensor(s25),
31 | torch.tensor(s13),
32 | torch.tensor(s7),
33 | ]
34 | )
35 |
36 | return ms_obj_size.cuda()
37 |
38 |
39 | def get_class_frequency():
40 | """
41 | NuScenes class frequency on Roddick Split
42 | gt_dir/name: semantic_maps_new_200x200
43 | """
44 |
45 | # Class frequency
46 | s200 = [
47 | 0.999,
48 | 0.489,
49 | 0.969,
50 | 0.294,
51 | 0.094,
52 | 0.080,
53 | 0.758,
54 | 0.0646,
55 | 0.070,
56 | 0.080,
57 | 0.313,
58 | 0.485,
59 | ]
60 |
61 | return torch.tensor(s200).cuda()
62 |
63 |
64 | def class_weight():
65 | """
66 | NuScenes class frequency on Roddick Split
67 | gt_dir/name: semantic_maps_new_200x200
68 | """
69 |
70 | # Class weights
71 | w = torch.tensor([1.31, 5.43, 1.00, 2.13, 8.04, 1.19, 1.75, 5.71])
72 |
73 | return w
74 |
75 |
76 | def class_flood_level():
77 | ms_flood = torch.tensor(
78 | [
79 | [
80 | 0.0356,
81 | 0.2656,
82 | 0.1070,
83 | 0.6454,
84 | 0.8,
85 | 0.75,
86 | 0.1832,
87 | 0.8,
88 | 0.75,
89 | 0.6633,
90 | 0.5591,
91 | ],
92 | [
93 | 0.0294,
94 | 0.2418,
95 | 0.0902,
96 | 0.5982,
97 | 0.75,
98 | 0.7,
99 | 0.1589,
100 | 0.8,
101 | 0.75,
102 | 0.6302,
103 | 0.5116,
104 | ],
105 | [
106 | 0.0247,
107 | 0.2363,
108 | 0.0787,
109 | 0.5765,
110 | 0.75,
111 | 0.75,
112 | 0.1416,
113 | 0.8,
114 | 0.75,
115 | 0.6317,
116 | 0.4982,
117 | ],
118 | ]
119 | ).cuda()
120 | return ms_flood
121 |
122 |
123 | class MultiTaskLoss(nn.Module):
124 | def __init__(self, l1_reduction="mean"):
125 | super().__init__()
126 |
127 | self.log_vars = nn.Parameter(torch.zeros(2))
128 | self.l1_reduction = l1_reduction
129 |
130 | def forward(self, preds, labels, bev, bev_pred):
131 | precision1 = torch.exp(-self.log_vars[0])
132 | precision2 = torch.exp(-self.log_vars[1])
133 |
134 | s1_loss = dice_loss_mean(preds[0], labels[0])
135 | s2_loss = dice_loss_mean(preds[1], labels[1])
136 | s4_loss = dice_loss_mean(preds[2], labels[2])
137 |
138 | dice_l = s1_loss + s2_loss + s4_loss
139 | bev_l = F.l1_loss(bev_pred, bev, reduction=self.l1_reduction)
140 |
141 | total_loss = (precision1 * dice_l + self.log_vars[0]) + (
142 | precision2 * bev_l + self.log_vars[1]
143 | )
144 |
145 | return (
146 | total_loss,
147 | float(self.log_vars[0]),
148 | float(self.log_vars[1]),
149 | float(dice_l),
150 | float(bev_l),
151 | )
152 |
153 |
154 | def lovasz_grad(gt_sorted):
155 | """
156 | Computes gradient of the Lovasz extension w.r.t sorted errors
157 | See Alg. 1 in paper
158 | """
159 | p = len(gt_sorted)
160 | gts = gt_sorted.sum()
161 | intersection = gts - gt_sorted.float().cumsum(0)
162 | union = gts + (1 - gt_sorted).float().cumsum(0)
163 | jaccard = 1.0 - intersection / union
164 | if p > 1: # cover 1-pixel case
165 | jaccard[1:p] = jaccard[1:p] - jaccard[0:-1]
166 | return jaccard
167 |
168 |
169 | def lovasz_hinge_flat(logits, labels):
170 | """
171 | Binary Lovasz hinge loss
172 | logits: [P] Variable, logits at each prediction (between -\infty and +\infty)
173 | labels: [P] Tensor, binary ground truth labels (0 or 1)
174 | ignore: label to ignore
175 | """
176 | logits = logits.view(-1)
177 | labels = labels.view(-1)
178 | signs = 2.0 * labels.float() - 1.0
179 | errors = 1.0 - logits * signs
180 | errors_sorted, perm = torch.sort(errors, dim=0, descending=True)
181 | perm = perm.data
182 | gt_sorted = labels[perm]
183 | grad = lovasz_grad(gt_sorted)
184 | loss = torch.dot(F.relu(errors_sorted), grad)
185 | return loss
186 |
187 |
188 | def basic_weighted_dice_loss_mean(pred, label, idx_scale=None, args=None):
189 | pred = torch.sigmoid(pred)
190 | label = label.float()
191 | intersection = 2 * pred * label
192 | union = pred + label
193 | iou = (intersection.float().sum(dim=0).sum(dim=-1).sum(dim=-1)) / (
194 | union.float().sum(dim=0).sum(dim=-1).sum(dim=-1) + 1e-5
195 | )
196 |
197 | # Set up class weighting
198 | weight = class_weight()
199 |
200 | loss_mean = (weight * (1 - iou)).mean()
201 |
202 | return loss_mean
203 |
204 |
205 | def weighted_dice_loss_mean(pred, label, idx_scale, args):
206 | pred = torch.sigmoid(pred)
207 | label = label.float()
208 | intersection = 2 * pred * label
209 | union = pred + label
210 | iou = (intersection.float().sum(dim=0).sum(dim=-1).sum(dim=-1)) / (
211 | union.float().sum(dim=0).sum(dim=-1).sum(dim=-1) + 1e-5
212 | )
213 |
214 | # Set up class weighting
215 | obj_size = class_obj_size()[idx_scale]
216 | class_freq = class_frequency()[idx_scale]
217 | weight = 1 / obj_size ** args.exp_os * 1 / class_freq ** args.exp_cf
218 |
219 | # Set large class' weights to one
220 | idxs_to_one = torch.tensor([0, 1, 2, 3, 6]).long()
221 | weight[idxs_to_one] = 1
222 |
223 | loss_mean = (weight * (1 - iou)).mean()
224 |
225 | return loss_mean
226 |
227 |
228 | def weighted_dice_loss_sum(pred, label, idx_scale, args):
229 | pred = torch.sigmoid(pred)
230 | label = label.float()
231 | intersection = 2 * pred * label
232 | union = pred + label
233 | iou = (intersection.float().sum(dim=0).sum(dim=-1).sum(dim=-1) + 1e-5) / (
234 | union.float().sum(dim=0).sum(dim=-1).sum(dim=-1) + 1e-5
235 | )
236 |
237 | # Set up class weighting
238 | obj_size = class_obj_size()[idx_scale]
239 | class_freq = class_frequency()[idx_scale]
240 | weight = 1 / obj_size ** args.exp_os * 1 / class_freq ** args.exp_cf
241 |
242 | loss_mean = (weight * (1 - iou)).sum()
243 |
244 | return loss_mean
245 |
246 |
247 | def dice_loss_sum(pred, label, scale_idx=None, args=None):
248 | pred = torch.sigmoid(pred)
249 | label = label.float()
250 | intersection = 2 * pred * label
251 | union = pred + label
252 | iou = (intersection.float().sum(dim=0).sum(dim=-1).sum(dim=-1)) / (
253 | union.float().sum(dim=0).sum(dim=-1).sum(dim=-1) + 1e-5
254 | )
255 |
256 | loss_mean = (1 - iou).sum()
257 |
258 | return loss_mean
259 |
260 |
261 | def dice_loss_mean(pred, label, vis_mask=None, scale_idx=None, args=None):
262 | pred = torch.sigmoid(pred)
263 | label = label.float()
264 | intersection = 2 * pred * label
265 | union = pred + label
266 | iou = (intersection.float().sum(dim=0).sum(dim=-1).sum(dim=-1)) / (
267 | union.float().sum(dim=0).sum(dim=-1).sum(dim=-1) + 1e-5
268 | )
269 |
270 | loss_mean = 1 - iou.mean()
271 |
272 | return loss_mean
273 |
274 |
275 | def dice_loss_mean_wo_sigmoid(pred, label, vis_mask=None, scale_idx=None, args=None):
276 | label = label.float()
277 | intersection = 2 * pred * label
278 | union = pred + label
279 | iou = (intersection.float().sum()) / (union.float().sum() + 1e-5)
280 |
281 | loss_mean = 1 - iou.mean()
282 |
283 | return loss_mean
284 |
285 |
286 | def dice_loss_mean_vis_only(pred, label, vis_mask, scale_idx=None, args=None):
287 | # Calculates the loss in the visible areas only
288 | pred = torch.sigmoid(pred) * vis_mask.float()
289 | label = label.float()
290 | intersection = 2 * pred * label
291 | union = pred + label
292 | iou = (intersection.float().sum(dim=0).sum(dim=-1).sum(dim=-1)) / (
293 | union.float().sum(dim=0).sum(dim=-1).sum(dim=-1) + 1e-5
294 | )
295 |
296 | loss_mean = 1 - iou.mean()
297 |
298 | return loss_mean
299 |
300 |
301 | def dice2_loss_mean(pred, label, scale_idx=None, args=None):
302 | pred = torch.sigmoid(pred)
303 | label = label.float()
304 | intersection = 2 * pred * label
305 | union = pred + label
306 | iou = (intersection.float().sum(dim=0).sum(dim=-1).sum(dim=-1) + 1.0) / (
307 | union.float().sum(dim=0).sum(dim=-1).sum(dim=-1) + 1.0
308 | )
309 |
310 | loss_mean = 1 - iou.mean()
311 |
312 | return loss_mean
313 |
314 |
315 | def dice_loss_mean_mix1(pred, label, scale_idx=None, args=None):
316 | pred = torch.sigmoid(pred)
317 | label = label.float()
318 | intersection = 2 * pred * label
319 | union = pred + label
320 |
321 | e_numerator = torch.zeros(pred.shape[1]) + 1e-5
322 | idxs_to_zero_out = torch.tensor([4, 5, 7, 8, 10]).long()
323 | e_numerator[idxs_to_zero_out] = 0
324 | e_numerator = e_numerator.cuda()
325 |
326 | iou = (intersection.float().sum(dim=0).sum(dim=-1).sum(dim=-1) + e_numerator) / (
327 | union.float().sum(dim=0).sum(dim=-1).sum(dim=-1) + 1e-5
328 | )
329 |
330 | loss_mean = 1 - iou.mean()
331 |
332 | return loss_mean
333 |
334 |
335 | def dice_loss_mean_mix2(pred, label, scale_idx=None, args=None):
336 | pred = torch.sigmoid(pred)
337 | label = label.float()
338 | intersection = 2 * pred * label
339 | union = pred + label
340 |
341 | e_numerator = torch.zeros(pred.shape[1]) + 1e-5
342 | idxs_to_zero_out = torch.tensor([4, 5, 7, 8]).long()
343 | e_numerator[idxs_to_zero_out] = 0
344 | e_numerator = e_numerator.cuda()
345 |
346 | iou = (intersection.float().sum(dim=0).sum(dim=-1).sum(dim=-1) + e_numerator) / (
347 | union.float().sum(dim=0).sum(dim=-1).sum(dim=-1) + 1e-5
348 | )
349 |
350 | loss_mean = 1 - iou.mean()
351 |
352 | return loss_mean
353 |
354 |
355 | def dice_loss_per_class(pred, labels, scale_idx=None):
356 | pred = torch.sigmoid(pred)
357 | labels = labels.float()
358 | intersection = 2 * pred * labels
359 | union = pred + labels
360 | iou = (intersection.float().sum(dim=0).sum(dim=-1).sum(dim=-1)) / (
361 | union.float().sum(dim=0).sum(dim=-1).sum(dim=-1) + 1e-5
362 | )
363 | return 1 - iou
364 |
365 |
366 | def dice_loss_per_class_infer(pred, labels, scale_idx=None):
367 | pred = torch.sigmoid(pred)
368 | labels = labels.float()
369 | intersection = (2 * pred * labels).float().sum(dim=0).sum(dim=-1).sum(dim=-1)
370 | union = (pred + labels).float().sum(dim=0).sum(dim=-1).sum(dim=-1)
371 | iou = (intersection) / (union + 1e-5)
372 | return 1 - iou, intersection, union
373 |
374 |
375 | def flooded_dice_loss_mean(dice_loss, scale_idx):
376 |
377 | # Flood level per class
378 | flood_level = class_flood_level()[scale_idx]
379 |
380 | flood_loss = (dice_loss - flood_level).abs() + flood_level
381 |
382 | return flood_loss.mean()
383 |
384 |
385 | def flooded_dice_loss_sum(dice_loss, scale_idx):
386 |
387 | # Flood level per class
388 | flood_level = class_flood_level()[scale_idx]
389 |
390 | flood_loss = (dice_loss - flood_level).abs() + flood_level
391 |
392 | return flood_loss.sum()
393 |
394 |
395 | def iou_loss(pred, labels):
396 | e = 1e-6
397 | pred = torch.sigmoid(pred)
398 | labels = labels.float()
399 | intersection = pred * labels
400 | union = (pred + labels) - (pred * labels)
401 | iou = intersection.sum() / (union.sum() + e)
402 | return 1.0 - iou
403 |
404 |
405 | def bce_plus_dice_loss(pred, labels, weights, alpha):
406 | bce = F.binary_cross_entropy_with_logits(
407 | input=pred, target=labels, weight=weights, reduction="mean"
408 | )
409 | dice = dice_loss_mean(pred, labels)
410 | total = (alpha[0] * bce + alpha[1] * dice).sum()
411 | return total
412 |
413 |
414 | def bce_plus_iou_loss(pred, labels, weights, alpha):
415 | bce = F.binary_cross_entropy_with_logits(
416 | input=pred, target=labels, weight=weights, reduction="mean"
417 | )
418 | iou = iou_loss(pred, labels)
419 | total = (alpha[0] * bce + alpha[1] * iou).sum()
420 | return total
421 |
422 |
423 | def class_balanced_loss(
424 | labels, logits, samples_per_cls, no_of_classes, loss_type, alpha=None
425 | ):
426 | """Compute the Class Balanced Loss between `logits` and the ground truth `labels`.
427 | Class Balanced Loss: ((1-beta)/(1-beta^n))*Loss(labels, logits)
428 | where Loss is one of the standard losses used for Neural Networks.
429 | Args:
430 | labels: A int tensor of size [batch].
431 | logits: A float tensor of size [batch, no_of_classes].
432 | samples_per_cls: A python list of size [no_of_classes].
433 | no_of_classes: total number of classes. int
434 | loss_type: string. One of "sigmoid", "focal", "softmax".
435 | beta: float. Hyperparameter for Class balanced loss.
436 | gamma: float. Hyperparameter for Focal loss.
437 | Returns:
438 | cb_loss: A float tensor representing class balanced loss
439 | """
440 | # effective_num = 1.0 - np.power(beta, samples_per_cls)
441 | # weights = (1.0 - beta) / np.array(effective_num)
442 | # weights = weights / np.sum(weights) * no_of_classes
443 |
444 | if loss_type == "bce":
445 | labels_one_hot = F.one_hot(labels.long(), no_of_classes).float()
446 | weights = 1 / np.array(samples_per_cls)
447 | weights = torch.tensor(weights).float().cuda()
448 | weights = weights.view(1, 1, 1, 1, -1)
449 | weights = (
450 | weights.repeat(
451 | labels_one_hot.shape[0],
452 | labels_one_hot.shape[1],
453 | labels_one_hot.shape[2],
454 | labels_one_hot.shape[3],
455 | 1,
456 | )
457 | * labels_one_hot
458 | )
459 | weights = weights.sum(-1)
460 | cb_loss = F.binary_cross_entropy_with_logits(
461 | input=logits, target=labels, weight=weights
462 | )
463 | elif loss_type == "iou":
464 | cb_loss = iou_loss(logits, labels)
465 | elif loss_type == "dice":
466 | cb_loss = dice_loss_mean(logits, labels)
467 | elif loss_type == "lovasz":
468 | cb_loss = lovasz_hinge_flat(logits, labels)
469 | # elif loss_type == "bce_iou":
470 | # cb_loss = bce_plus_iou_loss(logits, labels, weights, alpha)
471 | # elif loss_type == "bce_dice":
472 | # cb_loss = bce_plus_dice_loss(logits, labels, weights, alpha)
473 |
474 | return cb_loss
475 |
476 |
477 | def weighted_bce_loss(pred, label, vis_mask):
478 | # Apply visibility mask to predictions
479 | pred = pred * vis_mask.float()
480 |
481 | pred = torch.sigmoid(pred)
482 |
483 | label = label.float()
484 |
485 | eps = 1e-12
486 |
487 | # Reshape weights for broadcasting
488 | cf = get_class_frequency()[None, :, None, None]
489 | pos_weights = torch.sqrt(1 / cf)
490 | neg_weights = torch.sqrt(1 / (1 - cf))
491 |
492 | # labels * -log(sigmoid(logits)) +
493 | # (1 - labels) * -log(1 - sigmoid(logits))
494 |
495 | loss = pos_weights * label * -torch.log(pred + eps) + (1 - label) * -torch.log(
496 | 1 - pred + eps
497 | )
498 |
499 | return loss.sum(dim=-1).sum(dim=-1).sum(dim=0)
500 |
501 |
502 | def uncertainty_loss(pred, vis_mask):
503 | # Invert visibility mask
504 | vis_mask = 1 - vis_mask.float()
505 | eps = 1e-12
506 | pred = pred * vis_mask
507 | pred = torch.sigmoid(pred)
508 | loss = 1 - pred * torch.log(pred + eps)
509 |
510 | return loss.sum(dim=-1).sum(dim=-1).sum(dim=0)
511 |
512 |
513 | def aux_recognition_loss(pred, label):
514 | """
515 | BCE loss for class probabilities, multi-label classification
516 | """
517 | # pred = torch.sigmoid(pred)
518 | label = src.utils.count_classes_per_sample(label)
519 |
520 | assert pred.shape == label.shape
521 |
522 | class_freq = get_class_frequency()[None, :]
523 | weights = (1 / class_freq).repeat(len(pred), 1)
524 |
525 | loss = F.binary_cross_entropy_with_logits(
526 | pred, label, weight=weights, reduction="none"
527 | )
528 | return loss
529 |
530 |
531 | def focal_loss(pred, gt, alpha, gamma):
532 | BCE_loss = F.binary_cross_entropy_with_logits(
533 | pred, gt.float(), reduction="none"
534 | ).sum(-1)
535 | gt, _ = torch.max(gt, dim=-1)
536 | gt = gt.long()
537 | at = torch.tensor(alpha, device=gt.device).gather(0, gt.data.view(-1))
538 | pt = torch.exp(-BCE_loss)
539 | F_loss = at * (1 - pt) ** gamma * BCE_loss
540 | return F_loss
541 |
--------------------------------------------------------------------------------
/src/model/mocha2.py:
--------------------------------------------------------------------------------
1 | import torch
2 | from torch import nn
3 | import torch.nn.functional as F
4 |
5 |
6 | class Energy(nn.Module):
7 | def __init__(self, enc_dim=10, dec_dim=10, att_dim=10, init_r=-4):
8 | """
9 | [Modified Bahdahnau attention] from
10 | "Online and Linear-Time Attention by Enforcing Monotonic Alignment" (ICML 2017)
11 | http://arxiv.org/abs/1704.00784
12 | Used for Monotonic Attention and Chunk Attention
13 | """
14 | super().__init__()
15 | self.tanh = nn.Tanh()
16 | self.W = nn.Linear(enc_dim, att_dim, bias=False)
17 | self.V = nn.Linear(dec_dim, att_dim, bias=False)
18 | self.b = nn.Parameter(torch.Tensor(att_dim).normal_())
19 |
20 | self.v = nn.utils.weight_norm(nn.Linear(10, 1))
21 | self.v.weight_g.data = torch.Tensor([1 / att_dim]).sqrt()
22 |
23 | self.r = nn.Parameter(torch.Tensor([init_r]))
24 |
25 | def forward(self, encoder_outputs, decoder_h):
26 | """
27 | Args:
28 | encoder_outputs: [batch_size, sequence_length, enc_dim]
29 | decoder_h: [batch_size, dec_dim]
30 | Return:
31 | Energy [batch_size, sequence_length]
32 | """
33 | batch_size, sequence_length, enc_dim = encoder_outputs.size()
34 | encoder_outputs = encoder_outputs.view(-1, enc_dim)
35 | energy = self.tanh(
36 | self.W(encoder_outputs)
37 | + self.V(decoder_h).repeat(sequence_length, 1)
38 | + self.b
39 | )
40 | energy = self.v(energy).squeeze(-1) + self.r
41 |
42 | return energy.view(batch_size, sequence_length)
43 |
44 |
45 | class MonotonicAttention(nn.Module):
46 | def __init__(self):
47 | """
48 | [Monotonic Attention] from
49 | "Online and Linear-Time Attention by Enforcing Monotonic Alignment" (ICML 2017)
50 | http://arxiv.org/abs/1704.00784
51 | """
52 | super().__init__()
53 |
54 | self.monotonic_energy = Energy()
55 | self.sigmoid = nn.Sigmoid()
56 |
57 | def gaussian_noise(self, *size):
58 | """Additive gaussian nosie to encourage discreteness"""
59 | if torch.cuda.is_available():
60 | return torch.cuda.FloatTensor(*size).normal_()
61 | else:
62 | return torch.Tensor(*size).normal_()
63 |
64 | def safe_cumprod(self, x):
65 | """Numerically stable cumulative product by cumulative sum in log-space"""
66 | return torch.exp(
67 | torch.cumsum(torch.log(torch.clamp(x, min=1e-10, max=1)), dim=1)
68 | )
69 |
70 | def exclusive_cumprod(self, x):
71 | """Exclusive cumulative product [a, b, c] => [1, a, a * b]
72 | * TensorFlow: https://www.tensorflow.org/api_docs/python/tf/cumprod
73 | * PyTorch: https://discuss.pytorch.org/t/cumprod-exclusive-true-equivalences/2614
74 | """
75 | batch_size, sequence_length = x.size()
76 | if torch.cuda.is_available():
77 | one_x = torch.cat([torch.ones(batch_size, 1).cuda(), x], dim=1)[:, :-1]
78 | else:
79 | one_x = torch.cat([torch.ones(batch_size, 1), x], dim=1)[:, :-1]
80 | return torch.cumprod(one_x, dim=1)
81 |
82 | def soft(self, encoder_outputs, decoder_h, previous_alpha=None):
83 | """
84 | Soft monotonic attention (Train)
85 | Args:
86 | encoder_outputs [batch_size, sequence_length, enc_dim]
87 | decoder_h [batch_size, dec_dim]
88 | previous_alpha [batch_size, sequence_length]
89 | Return:
90 | alpha [batch_size, sequence_length]
91 | """
92 | batch_size, sequence_length, enc_dim = encoder_outputs.size()
93 |
94 | monotonic_energy = self.monotonic_energy(encoder_outputs, decoder_h)
95 | p_select = self.sigmoid(
96 | monotonic_energy + self.gaussian_noise(monotonic_energy.size())
97 | )
98 | cumprod_1_minus_p = self.safe_cumprod(1 - p_select)
99 |
100 | if previous_alpha is None:
101 | # First iteration => alpha = [1, 0, 0 ... 0]
102 | alpha = torch.zeros(batch_size, sequence_length)
103 | alpha[:, 0] = torch.ones(batch_size)
104 | if torch.cuda.is_available:
105 | alpha = alpha.cuda()
106 |
107 | else:
108 | alpha = (
109 | p_select
110 | * cumprod_1_minus_p
111 | * torch.cumsum(previous_alpha / cumprod_1_minus_p, dim=1)
112 | )
113 |
114 | return alpha
115 |
116 | def hard(self, encoder_outputs, decoder_h, previous_attention=None):
117 | """
118 | Hard monotonic attention (Test)
119 | Args:
120 | encoder_outputs [batch_size, sequence_length, enc_dim]
121 | decoder_h [batch_size, dec_dim]
122 | previous_attention [batch_size, sequence_length]
123 | Return:
124 | alpha [batch_size, sequence_length]
125 | """
126 | batch_size, sequence_length, enc_dim = encoder_outputs.size()
127 |
128 | if previous_attention is None:
129 | # First iteration => alpha = [1, 0, 0 ... 0]
130 | attention = torch.zeros(batch_size, sequence_length)
131 | attention[:, 0] = torch.ones(batch_size)
132 | if torch.cuda.is_available:
133 | attention = attention.cuda()
134 | else:
135 | # TODO: Linear Time Decoding
136 | # It's not clear if authors' TF implementation decodes in linear time.
137 | # https://github.com/craffel/mad/blob/master/example_decoder.py#L235
138 | # They calculate energies for whole encoder outputs
139 | # instead of scanning from previous attended encoder output.
140 | monotonic_energy = self.monotonic_energy(encoder_outputs, decoder_h)
141 |
142 | # Hard Sigmoid
143 | # Attend when monotonic energy is above threshold (Sigmoid > 0.5)
144 | above_threshold = (monotonic_energy > 0).float()
145 |
146 | p_select = above_threshold * torch.cumsum(previous_attention, dim=1)
147 | attention = p_select * self.exclusive_cumprod(1 - p_select)
148 |
149 | # Not attended => attend at last encoder output
150 | # Assume that encoder outputs are not padded
151 | attended = attention.sum(dim=1)
152 | for batch_i in range(batch_size):
153 | if not attended[batch_i]:
154 | attention[batch_i, -1] = 1
155 |
156 | # Ex)
157 | # p_select = [0, 0, 0, 1, 1, 0, 1, 1]
158 | # 1 - p_select = [1, 1, 1, 0, 0, 1, 0, 0]
159 | # exclusive_cumprod(1 - p_select) = [1, 1, 1, 1, 0, 0, 0, 0]
160 | # attention: product of above = [0, 0, 0, 1, 0, 0, 0, 0]
161 | return attention
162 |
163 |
164 | class MoChA(MonotonicAttention):
165 | def __init__(self, chunk_size=3):
166 | """
167 | [Monotonic Chunkwise Attention] from
168 | "Monotonic Chunkwise Attention" (ICLR 2018)
169 | https://openreview.net/forum?id=Hko85plCW
170 | """
171 | super().__init__()
172 | self.chunk_size = chunk_size
173 | self.chunk_energy = Energy()
174 | self.softmax = nn.Softmax(dim=1)
175 |
176 | def moving_sum(self, x, back, forward):
177 | """Parallel moving sum with 1D Convolution"""
178 | # Pad window before applying convolution
179 | # [batch_size, back + sequence_length + forward]
180 | x_padded = F.pad(x, pad=[back, forward])
181 |
182 | # Fake channel dimension for conv1d
183 | # [batch_size, 1, back + sequence_length + forward]
184 | x_padded = x_padded.unsqueeze(1)
185 |
186 | # Apply conv1d with filter of all ones for moving sum
187 | filters = torch.ones(1, 1, back + forward + 1)
188 | if torch.cuda.is_available():
189 | filters = filters.cuda()
190 | x_sum = F.conv1d(x_padded, filters)
191 |
192 | # Remove fake channel dimension
193 | # [batch_size, sequence_length]
194 | return x_sum.squeeze(1)
195 |
196 | def chunkwise_attention_soft(self, alpha, u):
197 | """
198 | Args:
199 | alpha [batch_size, sequence_length]: emission probability in monotonic attention
200 | u [batch_size, sequence_length]: chunk energy
201 | chunk_size (int): window size of chunk
202 | Return
203 | beta [batch_size, sequence_length]: MoChA weights
204 | """
205 |
206 | # Numerical stability
207 | # Divide by same exponent => doesn't affect softmax
208 | u -= torch.max(u, dim=1, keepdim=True)[0]
209 | exp_u = torch.exp(u)
210 | # Limit range of logit
211 | exp_u = torch.clamp(exp_u, min=1e-5)
212 |
213 | # Moving sum:
214 | # Zero-pad (chunk size - 1) on the left + 1D conv with filters of 1s.
215 | # [batch_size, sequence_length]
216 | denominators = self.moving_sum(exp_u, back=self.chunk_size - 1, forward=0)
217 |
218 | # Compute beta (MoChA weights)
219 | beta = exp_u * self.moving_sum(
220 | alpha / denominators, back=0, forward=self.chunk_size - 1
221 | )
222 | return beta
223 |
224 | def chunkwise_attention_hard(self, monotonic_attention, chunk_energy):
225 | """
226 | Mask non-attended area with '-inf'
227 | Args:
228 | monotonic_attention [batch_size, sequence_length]
229 | chunk_energy [batch_size, sequence_length]
230 | Return:
231 | masked_energy [batch_size, sequence_length]
232 | """
233 | batch_size, sequence_length = monotonic_attention.size()
234 |
235 | # [batch_size]
236 | attended_indices = monotonic_attention.nonzero().cpu().data[:, 1].tolist()
237 |
238 | i = [[], []]
239 | total_i = 0
240 | for batch_i, attended_idx in enumerate(attended_indices):
241 | for window in range(self.chunk_size):
242 | if attended_idx - window >= 0:
243 | i[0].append(batch_i)
244 | i[1].append(attended_idx - window)
245 | total_i += 1
246 | i = torch.LongTensor(i)
247 | v = torch.FloatTensor([1] * total_i)
248 | mask = torch.sparse.FloatTensor(i, v, monotonic_attention.size())
249 | mask = ~mask.to_dense().cuda().byte()
250 |
251 | # mask '-inf' energy before softmax
252 | masked_energy = chunk_energy.masked_fill_(mask, -float("inf"))
253 | return masked_energy
254 |
255 | def soft(self, encoder_outputs, decoder_h, previous_alpha=None):
256 | """
257 | Soft monotonic chunkwise attention (Train)
258 | Args:
259 | encoder_outputs [batch_size, sequence_length, enc_dim]
260 | decoder_h [batch_size, dec_dim]
261 | previous_alpha [batch_size, sequence_length]
262 | Return:
263 | alpha [batch_size, sequence_length]
264 | beta [batch_size, sequence_length]
265 | """
266 | alpha = super().soft(encoder_outputs, decoder_h, previous_alpha)
267 | chunk_energy = self.chunk_energy(encoder_outputs, decoder_h)
268 | beta = self.chunkwise_attention_soft(alpha, chunk_energy)
269 | return alpha, beta
270 |
271 | def hard(self, encoder_outputs, decoder_h, previous_attention=None):
272 | """
273 | Hard monotonic chunkwise attention (Test)
274 | Args:
275 | encoder_outputs [batch_size, sequence_length, enc_dim]
276 | decoder_h [batch_size, dec_dim]
277 | previous_attention [batch_size, sequence_length]
278 | Return:
279 | monotonic_attention [batch_size, sequence_length]: hard alpha
280 | chunkwise_attention [batch_size, sequence_length]: hard beta
281 | """
282 | # hard attention (one-hot)
283 | # [batch_size, sequence_length]
284 | monotonic_attention = super().hard(
285 | encoder_outputs, decoder_h, previous_attention
286 | )
287 | chunk_energy = self.chunk_energy(encoder_outputs, decoder_h)
288 | masked_energy = self.chunkwise_attention_hard(monotonic_attention, chunk_energy)
289 | chunkwise_attention = self.softmax(masked_energy)
290 | return monotonic_attention, chunkwise_attention
291 |
--------------------------------------------------------------------------------
/src/model/monotonic_attention.py:
--------------------------------------------------------------------------------
1 | import math
2 | import numpy as np
3 | import random
4 | import torch
5 | import torch.nn as nn
6 | import torch.nn.functional as F
7 |
8 |
9 | class MonotonicAttention(nn.Module):
10 | """
11 | Monotonic Multihead Attention with Infinite Lookback
12 | """
13 |
14 | def __init__(
15 | self,
16 | embed_dim,
17 | num_heads,
18 | kdim=None,
19 | vdim=None,
20 | dropout=0.0,
21 | bias=True,
22 | energy_bias=True,
23 | ):
24 | self.embed_dim = embed_dim
25 | self.kdim = kdim if kdim is not None else embed_dim
26 | self.vdim = vdim if vdim is not None else embed_dim
27 | self.qkv_same_dim = self.kdim == embed_dim and self.vdim == embed_dim
28 |
29 | self.k_proj_mono = nn.Linear(embed_dim, embed_dim, bias=bias)
30 | self.q_proj_mono = nn.Linear(embed_dim, embed_dim, bias=bias)
31 | self.k_proj_soft = nn.Linear(embed_dim, embed_dim, bias=bias)
32 | self.q_proj_soft = nn.Linear(embed_dim, embed_dim, bias=bias)
33 | self.v_proj = nn.Linear(self.vdim, embed_dim, bias=bias)
34 |
35 | self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
36 |
37 | self.num_heads = num_heads
38 | self.head_dim = embed_dim // num_heads
39 | assert (
40 | self.head_dim * num_heads == self.embed_dim
41 | ), "embed_dim must be divisible by num_heads"
42 | self.scaling = self.head_dim ** -0.5
43 |
44 | self.eps = 1e-6
45 |
46 | self.noise_mean = 0.0
47 | self.noise_var = 1.0
48 |
49 | self.energy_bias_init = -2.0
50 | self.energy_bias = (
51 | nn.Parameter(self.energy_bias_init * torch.ones([1]))
52 | if energy_bias is True
53 | else 0
54 | )
55 |
56 | self.reset_parameters()
57 |
58 | self.k_in_proj = {"monotonic": self.k_proj_mono, "soft": self.k_proj_soft}
59 | self.q_in_proj = {"monotonic": self.q_proj_mono, "soft": self.q_proj_soft}
60 | self.v_in_proj = {"output": self.v_proj}
61 |
62 | def reset_parameters(self):
63 | if self.qkv_same_dim:
64 | # Empirically observed the convergence to be much better with
65 | # the scaled initialization
66 | nn.init.xavier_uniform_(self.k_proj_mono.weight, gain=1 / math.sqrt(2))
67 | nn.init.xavier_uniform_(self.k_proj_soft.weight, gain=1 / math.sqrt(2))
68 | nn.init.xavier_uniform_(self.q_proj_mono.weight, gain=1 / math.sqrt(2))
69 | nn.init.xavier_uniform_(self.q_proj_soft.weight, gain=1 / math.sqrt(2))
70 | nn.init.xavier_uniform_(self.v_proj.weight, gain=1 / math.sqrt(2))
71 |
72 | else:
73 | nn.init.xavier_uniform_(self.k_proj_mono.weight)
74 | nn.init.xavier_uniform_(self.k_proj_soft.weight)
75 | nn.init.xavier_uniform_(self.q_proj_mono.weight)
76 | nn.init.xavier_uniform_(self.q_proj_soft.weight)
77 | nn.init.xavier_uniform_(self.v_proj.weight)
78 |
79 | nn.init.xavier_uniform_(self.out_proj.weight)
80 | if self.out_proj.bias is not None:
81 | nn.init.constant_(self.out_proj.bias, 0.0)
82 |
83 | def input_projections(self, query, key, value, name):
84 | """
85 | Prepare inputs for multihead attention
86 | ============================================================
87 | Expected input size
88 | query: tgt_len, bsz, embed_dim
89 | key: src_len, bsz, embed_dim
90 | value: src_len, bsz, embed_dim
91 | name: monotonic or soft
92 | """
93 |
94 | if query is not None:
95 | bsz = query.size(1)
96 | q = self.q_in_proj(query)
97 | q *= self.scaling
98 | q = (
99 | q.contiguous()
100 | .view(-1, bsz * self.num_heads, self.head_dim)
101 | .transpose(0, 1)
102 | )
103 | else:
104 | q = None
105 |
106 | if key is not None:
107 | bsz = key.size(1)
108 | k = self.k_in_proj(key)
109 | k = (
110 | k.contiguous()
111 | .view(-1, bsz * self.num_heads, self.head_dim)
112 | .transpose(0, 1)
113 | )
114 | else:
115 | k = None
116 |
117 | if value is not None:
118 | bsz = value.size(1)
119 | v = self.v_proj(value)
120 | v = (
121 | v.contiguous()
122 | .view(-1, bsz * self.num_heads, self.head_dim)
123 | .transpose(0, 1)
124 | )
125 | else:
126 | v = None
127 |
128 | return q, k, v
129 |
130 | def p_choose(self, query, key, key_padding_mask=None):
131 | """
132 | Calculating step wise prob for reading and writing
133 | 1 to read, 0 to write
134 | ============================================================
135 | Expected input size
136 | query: bsz, tgt_len, embed_dim
137 | key: bsz, src_len, embed_dim
138 | value: bsz, src_len, embed_dim
139 | key_padding_mask: bsz, src_len
140 | attn_mask: bsz, src_len
141 | query: bsz, tgt_len, embed_dim
142 | """
143 |
144 | # prepare inputs
145 | q_proj, k_proj, _ = self.input_projections(query, key, None, "monotonic")
146 |
147 | # attention energy
148 | attn_energy = self.attn_energy(q_proj, k_proj, key_padding_mask)
149 |
150 | noise = 0
151 |
152 | if self.training:
153 | # add noise here to encourage discretness
154 | noise = (
155 | torch.normal(self.noise_mean, self.noise_var, attn_energy.size())
156 | .type_as(attn_energy)
157 | .to(attn_energy.device)
158 | )
159 |
160 | p_choose = torch.sigmoid(attn_energy + noise)
161 | _, _, tgt_len, src_len = p_choose.size()
162 |
163 | # p_choose: bsz * self.num_heads, tgt_len, src_len
164 | return p_choose.view(-1, tgt_len, src_len)
165 |
166 | def attn_energy(self, q_proj, k_proj, key_padding_mask=None):
167 | """
168 | Calculating monotonic energies
169 | ============================================================
170 | Expected input size
171 | q_proj: bsz * num_heads, tgt_len, self.head_dim
172 | k_proj: bsz * num_heads, src_len, self.head_dim
173 | key_padding_mask: bsz, src_len
174 | attn_mask: tgt_len, src_len
175 | """
176 | bsz, tgt_len, embed_dim = q_proj.size()
177 | bsz = bsz // self.num_heads
178 | src_len = k_proj.size(1)
179 |
180 | attn_energy = torch.bmm(q_proj, k_proj.transpose(1, 2)) + self.energy_bias
181 |
182 | attn_energy = attn_energy.view(bsz, self.num_heads, tgt_len, src_len)
183 |
184 | if key_padding_mask is not None:
185 | attn_energy = attn_energy.masked_fill(
186 | key_padding_mask.unsqueeze(1).unsqueeze(2).bool(),
187 | float("-inf"),
188 | )
189 |
190 | return attn_energy
191 |
192 | def expected_alignment_train(self, p_choose, key_padding_mask):
193 | """
194 | Calculating expected alignment for MMA
195 | Mask is not need because p_choose will be 0 if masked
196 | q_ij = (1 − p_{ij−1})q_{ij−1} + a+{i−1j}
197 | a_ij = p_ij q_ij
198 | parellel solution:
199 | ai = p_i * cumprod(1 − pi) * cumsum(a_i / cumprod(1 − pi))
200 | ============================================================
201 | Expected input size
202 | p_choose: bsz * num_heads, tgt_len, src_len
203 | """
204 |
205 | # p_choose: bsz * num_heads, tgt_len, src_len
206 | bsz_num_heads, tgt_len, src_len = p_choose.size()
207 |
208 | # cumprod_1mp : bsz * num_heads, tgt_len, src_len
209 | cumprod_1mp = exclusive_cumprod(1 - p_choose, dim=2, eps=self.eps)
210 | cumprod_1mp_clamp = torch.clamp(cumprod_1mp, self.eps, 1.0)
211 |
212 | init_attention = p_choose.new_zeros([bsz_num_heads, 1, src_len])
213 | init_attention[:, :, 0] = 1.0
214 |
215 | previous_attn = [init_attention]
216 |
217 | for i in range(tgt_len):
218 | # p_choose: bsz * num_heads, tgt_len, src_len
219 | # cumprod_1mp_clamp : bsz * num_heads, tgt_len, src_len
220 | # previous_attn[i]: bsz * num_heads, 1, src_len
221 | # alpha_i: bsz * num_heads, src_len
222 | alpha_i = (
223 | p_choose[:, i]
224 | * cumprod_1mp[:, i]
225 | * torch.cumsum(previous_attn[i][:, 0] / cumprod_1mp_clamp[:, i], dim=1)
226 | ).clamp(0, 1.0)
227 | previous_attn.append(alpha_i.unsqueeze(1))
228 |
229 | # alpha: bsz * num_heads, tgt_len, src_len
230 | alpha = torch.cat(previous_attn[1:], dim=1)
231 |
232 | if self.mass_preservation:
233 | # Last token has the residual probabilities
234 | alpha[:, :, -1] = 1 - alpha[:, :, :-1].sum(dim=-1).clamp(0.0, 1.0)
235 |
236 | assert not torch.isnan(alpha).any(), "NaN detected in alpha."
237 |
238 | return alpha
239 |
240 | def expected_attention(
241 | self, alpha, query, key, value, key_padding_mask, incremental_state
242 | ):
243 | # monotonic attention, we will calculate milk here
244 | bsz_x_num_heads, tgt_len, src_len = alpha.size()
245 | bsz = int(bsz_x_num_heads / self.num_heads)
246 |
247 | q, k, _ = self.input_projections(query, key, None, "soft")
248 | soft_energy = self.attn_energy(q, k, key_padding_mask)
249 |
250 | assert list(soft_energy.size()) == [bsz, self.num_heads, tgt_len, src_len]
251 |
252 | soft_energy = soft_energy.view(bsz * self.num_heads, tgt_len, src_len)
253 |
254 | if incremental_state is not None:
255 | monotonic_cache = self._get_monotonic_buffer(incremental_state)
256 | monotonic_step = monotonic_cache["step"] + 1
257 | step_offset = 0
258 | if key_padding_mask is not None:
259 | if key_padding_mask[:, 0].any():
260 | # left_pad_source = True:
261 | step_offset = key_padding_mask.sum(dim=-1, keepdim=True)
262 | monotonic_step += step_offset
263 | mask = lengths_to_mask(
264 | monotonic_step.view(-1), soft_energy.size(2), 1
265 | ).unsqueeze(1)
266 |
267 | soft_energy = soft_energy.masked_fill(~mask.bool(), float("-inf"))
268 | soft_energy = soft_energy - soft_energy.max(dim=2, keepdim=True)[0]
269 | exp_soft_energy = torch.exp(soft_energy)
270 | exp_soft_energy_sum = exp_soft_energy.sum(dim=2)
271 | beta = exp_soft_energy / exp_soft_energy_sum.unsqueeze(2)
272 |
273 | else:
274 | # bsz * num_heads, tgt_len, src_len
275 | soft_energy = soft_energy - soft_energy.max(dim=2, keepdim=True)[0]
276 | exp_soft_energy = torch.exp(soft_energy)
277 | exp_soft_energy_cumsum = torch.cumsum(exp_soft_energy, dim=2)
278 |
279 | if key_padding_mask is not None:
280 | if key_padding_mask.any():
281 | exp_soft_energy_cumsum = (
282 | exp_soft_energy_cumsum.view(
283 | -1, self.num_heads, tgt_len, src_len
284 | )
285 | .masked_fill(
286 | key_padding_mask.unsqueeze(1).unsqueeze(1), self.eps
287 | )
288 | .view(-1, tgt_len, src_len)
289 | )
290 |
291 | inner_items = alpha / exp_soft_energy_cumsum
292 |
293 | beta = exp_soft_energy * torch.cumsum(
294 | inner_items.flip(dims=[2]), dim=2
295 | ).flip(dims=[2])
296 |
297 | beta = self.dropout_module(beta)
298 |
299 | assert not torch.isnan(beta).any(), "NaN detected in beta."
300 |
301 | return beta
302 |
303 | def forward(
304 | self,
305 | query,
306 | key,
307 | value,
308 | key_padding_mask=None,
309 | incremental_state=None,
310 | *args,
311 | **kwargs,
312 | ):
313 |
314 | tgt_len, bsz, embed_dim = query.size()
315 | src_len = value.size(0)
316 |
317 | # stepwise prob
318 | # p_choose: bsz * self.num_heads, tgt_len, src_len
319 | p_choose = self.p_choose(query, key, key_padding_mask)
320 |
321 | # expected alignment alpha
322 | # bsz * self.num_heads, tgt_len, src_len
323 |
324 | alpha = self.expected_alignment_train(p_choose, key_padding_mask)
325 |
326 | # expected attention beta
327 | # bsz * self.num_heads, tgt_len, src_len
328 | beta = self.expected_attention(
329 | alpha, query, key, value, key_padding_mask, incremental_state
330 | )
331 |
332 | attn_weights = beta
333 |
334 | _, _, v_proj = self.input_projections(None, None, value, "output")
335 | attn = torch.bmm(attn_weights.type_as(v_proj), v_proj)
336 |
337 | attn = attn.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)
338 |
339 | attn = self.out_proj(attn)
340 |
341 | beta = beta.view(bsz, self.num_heads, tgt_len, src_len)
342 | alpha = alpha.view(bsz, self.num_heads, tgt_len, src_len)
343 | p_choose = p_choose.view(bsz, self.num_heads, tgt_len, src_len)
344 |
345 | return attn, alpha, beta, p_choose
346 |
347 |
348 | def exclusive_cumprod(tensor, dim: int, eps: float = 1e-10):
349 | """
350 | Implementing exclusive cumprod.
351 | There is cumprod in pytorch, however there is no exclusive mode.
352 | cumprod(x) = [x1, x1x2, x2x3x4, ..., prod_{i=1}^n x_i]
353 | exclusive means cumprod(x) = [1, x1, x1x2, x1x2x3, ..., prod_{i=1}^{n-1} x_i]
354 | """
355 | tensor_size = list(tensor.size())
356 | tensor_size[dim] = 1
357 | return_tensor = safe_cumprod(
358 | torch.cat([torch.ones(tensor_size).type_as(tensor), tensor], dim=dim),
359 | dim=dim,
360 | eps=eps,
361 | )
362 |
363 | if dim == 0:
364 | return return_tensor[:-1]
365 | elif dim == 1:
366 | return return_tensor[:, :-1]
367 | elif dim == 2:
368 | return return_tensor[:, :, :-1]
369 | else:
370 | raise RuntimeError("Cumprod on dimension 3 and more is not implemented")
371 |
372 |
373 | def safe_cumprod(tensor, dim: int, eps: float = 1e-10):
374 | """
375 | An implementation of cumprod to prevent precision issue.
376 | cumprod(x)
377 | = [x1, x1x2, x1x2x3, ....]
378 | = [exp(log(x1)), exp(log(x1) + log(x2)), exp(log(x1) + log(x2) + log(x3)), ...]
379 | = exp(cumsum(log(x)))
380 | """
381 |
382 | if (tensor + eps < 0).any().item():
383 | raise RuntimeError(
384 | "Safe cumprod can only take non-negative tensors as input."
385 | "Consider use torch.cumprod if you want to calculate negative values."
386 | )
387 |
388 | log_tensor = torch.log(tensor + eps)
389 | cumsum_log_tensor = torch.cumsum(log_tensor, dim)
390 | exp_cumsum_log_tensor = torch.exp(cumsum_log_tensor)
391 | return exp_cumsum_log_tensor
392 |
393 |
394 | def lengths_to_mask(lengths, max_len: int, dim: int = 0, negative_mask: bool = False):
395 | """
396 | Convert a tensor of lengths to mask
397 | For example, lengths = [[2, 3, 4]], max_len = 5
398 | mask =
399 | [[1, 1, 1],
400 | [1, 1, 1],
401 | [0, 1, 1],
402 | [0, 0, 1],
403 | [0, 0, 0]]
404 | """
405 | assert len(lengths.size()) <= 2
406 | if len(lengths) == 2:
407 | if dim == 1:
408 | lengths = lengths.t()
409 | lengths = lengths
410 | else:
411 | lengths = lengths.unsqueeze(1)
412 |
413 | # lengths : batch_size, 1
414 | lengths = lengths.view(-1, 1)
415 |
416 | batch_size = lengths.size(0)
417 | # batch_size, max_len
418 | mask = torch.arange(max_len).expand(batch_size, max_len).type_as(lengths) < lengths
419 |
420 | if negative_mask:
421 | mask = ~mask
422 |
423 | if dim == 0:
424 | # max_len, batch_size
425 | mask = mask.t()
426 |
427 | return mask
428 |
429 |
430 | def moving_sum(x, start_idx: int, end_idx: int):
431 | """
432 | From MONOTONIC CHUNKWISE ATTENTION
433 | https://arxiv.org/pdf/1712.05382.pdf
434 | Equation (18)
435 | x = [x_1, x_2, ..., x_N]
436 | MovingSum(x, start_idx, end_idx)_n = Sigma_{m=n−(start_idx−1)}^{n+end_idx-1} x_m
437 | for n in {1, 2, 3, ..., N}
438 | x : src_len, batch_size
439 | start_idx : start idx
440 | end_idx : end idx
441 | Example
442 | src_len = 5
443 | batch_size = 3
444 | x =
445 | [[ 0, 5, 10],
446 | [ 1, 6, 11],
447 | [ 2, 7, 12],
448 | [ 3, 8, 13],
449 | [ 4, 9, 14]]
450 | MovingSum(x, 3, 1) =
451 | [[ 0, 5, 10],
452 | [ 1, 11, 21],
453 | [ 3, 18, 33],
454 | [ 6, 21, 36],
455 | [ 9, 24, 39]]
456 | MovingSum(x, 1, 3) =
457 | [[ 3, 18, 33],
458 | [ 6, 21, 36],
459 | [ 9, 24, 39],
460 | [ 7, 17, 27],
461 | [ 4, 9, 14]]
462 | """
463 | assert start_idx > 0 and end_idx > 0
464 | assert len(x.size()) == 2
465 | src_len, batch_size = x.size()
466 | # batch_size, 1, src_len
467 | x = x.t().unsqueeze(1)
468 | # batch_size, 1, src_len
469 | moving_sum_weight = x.new_ones([1, 1, end_idx + start_idx - 1])
470 |
471 | moving_sum = (
472 | torch.nn.functional.conv1d(
473 | x, moving_sum_weight, padding=start_idx + end_idx - 1
474 | )
475 | .squeeze(1)
476 | .t()
477 | )
478 | moving_sum = moving_sum[end_idx:-start_idx]
479 |
480 | assert src_len == moving_sum.size(0)
481 | assert batch_size == moving_sum.size(1)
482 |
483 | return moving_sum
484 |
--------------------------------------------------------------------------------
/src/model/monotonic_attention00.py:
--------------------------------------------------------------------------------
1 | import math
2 | import numpy as np
3 | import random
4 | import torch
5 | import torch.nn as nn
6 | import torch.nn.functional as F
7 |
8 |
9 | class MonotonicEnergy(nn.Module):
10 | def __init__(
11 | self,
12 | kdim,
13 | qdim,
14 | adim,
15 | atype,
16 | n_heads,
17 | init_r,
18 | bias=True,
19 | param_init="",
20 | conv1d=False,
21 | conv_kernel_size=5,
22 | ):
23 | """Energy function for the monotonic attenion.
24 | Args:
25 | kdim (int): dimension of key
26 | qdim (int): dimension of quary
27 | adim (int): dimension of attention space
28 | atype (str): type of attention mechanism
29 | n_heads (int): number of monotonic attention heads
30 | init_r (int): initial value for offset r
31 | bias (bool): use bias term in linear layers
32 | param_init (str): parameter initialization method
33 | conv1d (bool): use 1D causal convolution for energy calculation
34 | conv_kernel_size (int): kernel size for 1D convolution
35 | """
36 | super().__init__()
37 |
38 | assert conv_kernel_size % 2 == 1, "Kernel size should be odd for 'same' conv."
39 | self.key = None
40 | self.mask = None
41 |
42 | self.atype = atype
43 | assert adim % n_heads == 0
44 | self.d_k = adim // n_heads
45 | self.n_heads = n_heads
46 | self.scale = math.sqrt(adim)
47 |
48 | if atype == "add":
49 | self.w_key = nn.Linear(kdim, adim)
50 | self.v = nn.Linear(adim, n_heads, bias=False)
51 | self.w_query = nn.Linear(qdim, adim, bias=False)
52 | elif atype == "scaled_dot":
53 | self.w_key = nn.Linear(kdim, adim, bias=bias)
54 | self.w_query = nn.Linear(qdim, adim, bias=bias)
55 | else:
56 | raise NotImplementedError(atype)
57 |
58 | self.r = nn.Parameter(torch.Tensor([init_r]))
59 | logger.info("init_r is initialized with %d" % init_r)
60 |
61 | self.conv1d = None
62 | if conv1d:
63 | self.conv1d = CausalConv1d(
64 | in_channels=kdim,
65 | out_channels=kdim,
66 | kernel_size=conv_kernel_size,
67 | param_init=param_init,
68 | )
69 | # padding=(conv_kernel_size - 1) // 2
70 |
71 | if atype == "add":
72 | self.v = nn.utils.weight_norm(self.v, name="weight", dim=0)
73 | # initialization
74 | self.v.weight_g.data = torch.Tensor([1 / adim]).sqrt()
75 | elif atype == "scaled_dot":
76 | if param_init == "xavier_uniform":
77 | self.reset_parameters(bias)
78 |
79 | def reset_parameters(self, bias):
80 | """Initialize parameters with Xavier uniform distribution."""
81 | logger.info(
82 | "===== Initialize %s with Xavier uniform distribution ====="
83 | % self.__class__.__name__
84 | )
85 | # NOTE: see https://github.com/pytorch/fairseq/blob/master/fairseq/modules/multihead_attention.py
86 | nn.init.xavier_uniform_(self.w_key.weight, gain=1 / math.sqrt(2))
87 | nn.init.xavier_uniform_(self.w_query.weight, gain=1 / math.sqrt(2))
88 | if bias:
89 | nn.init.constant_(self.w_key.bias, 0.0)
90 | nn.init.constant_(self.w_query.bias, 0.0)
91 |
92 | def reset(self):
93 | self.key = None
94 | self.mask = None
95 |
96 | def forward(self, key, query, mask, cache=False, boundary_leftmost=0):
97 | """Compute monotonic energy.
98 | Args:
99 | key (FloatTensor): `[B, klen, kdim]`
100 | query (FloatTensor): `[B, qlen, qdim]`
101 | mask (ByteTensor): `[B, qlen, klen]`
102 | cache (bool): cache key and mask
103 | Returns:
104 | e (FloatTensor): `[B, H_ma, qlen, klen]`
105 | """
106 | bs, klen, kdim = key.size()
107 | qlen = query.size(1)
108 |
109 | # Pre-computation of encoder-side features for computing scores
110 | if self.key is None or not cache:
111 | # 1d conv
112 | if self.conv1d is not None:
113 | key = torch.relu(self.conv1d(key))
114 | key = self.w_key(key).view(bs, -1, self.n_heads, self.d_k)
115 | self.key = key.transpose(2, 1).contiguous() # `[B, H_ma, klen, d_k]`
116 | self.mask = mask
117 | if mask is not None:
118 | self.mask = self.mask.unsqueeze(1).repeat(
119 | [1, self.n_heads, 1, 1]
120 | ) # `[B, H_ma, qlen, klen]`
121 | assert self.mask.size() == (bs, self.n_heads, qlen, klen), (
122 | self.mask.size(),
123 | (bs, self.n_heads, qlen, klen),
124 | )
125 |
126 | query = self.w_query(query).view(bs, -1, self.n_heads, self.d_k)
127 | query = query.transpose(2, 1).contiguous() # `[B, H_ma, qlen, d_k]`
128 | m = self.mask
129 |
130 | if self.atype == "add":
131 | k = self.key.unsqueeze(2) # `[B, H_ma, 1, klen, d_k]`
132 | # Truncate encoder memories
133 | if boundary_leftmost > 0:
134 | k = k[:, :, :, boundary_leftmost:]
135 | klen = k.size(3)
136 | if m is not None:
137 | m = m[:, :, :, boundary_leftmost:]
138 | e = torch.relu(k + query.unsqueeze(3)) # `[B, H_ma, qlen, klen, d_k]`
139 | e = e.permute(0, 2, 3, 1, 4).contiguous().view(bs, qlen, klen, -1)
140 | e = self.v(e).permute(0, 3, 1, 2) # `[B, qlen, klen, H_ma]`
141 | elif self.atype == "scaled_dot":
142 | k = self.key.transpose(3, 2)
143 | e = torch.matmul(query, k) / self.scale
144 |
145 | if self.r is not None:
146 | e = e + self.r
147 | if m is not None:
148 | NEG_INF = float(np.finfo(torch.tensor(0, dtype=e.dtype).numpy().dtype).min)
149 | e = e.masked_fill_(m == 0, NEG_INF)
150 | assert e.size() == (bs, self.n_heads, qlen, klen), (
151 | e.size(),
152 | (bs, self.n_heads, qlen, klen),
153 | )
154 | return e
155 |
156 |
157 | class ChunkEnergy(nn.Module):
158 | def __init__(self, kdim, qdim, adim, atype, n_heads=1, bias=True, param_init=""):
159 | """Energy function for the chunkwise attention.
160 | Args:
161 | kdim (int): dimension of key
162 | qdim (int): dimension of quary
163 | adim (int): dimension of attention space
164 | atype (str): type of attention mechanism
165 | n_heads (int): number of chunkwise attention heads
166 | bias (bool): use bias term in linear layers
167 | param_init (str): parameter initialization method
168 | """
169 | super().__init__()
170 |
171 | self.key = None
172 | self.mask = None
173 |
174 | self.atype = atype
175 | assert adim % n_heads == 0
176 | self.d_k = adim // n_heads
177 | self.n_heads = n_heads
178 | self.scale = math.sqrt(adim)
179 |
180 | if atype == "add":
181 | self.w_key = nn.Linear(kdim, adim)
182 | self.w_query = nn.Linear(qdim, adim, bias=False)
183 | self.v = nn.Linear(adim, n_heads, bias=False)
184 | elif atype == "scaled_dot":
185 | self.w_key = nn.Linear(kdim, adim, bias=bias)
186 | self.w_query = nn.Linear(qdim, adim, bias=bias)
187 | if param_init == "xavier_uniform":
188 | self.reset_parameters(bias)
189 | else:
190 | raise NotImplementedError(atype)
191 |
192 | def reset_parameters(self, bias):
193 | """Initialize parameters with Xavier uniform distribution."""
194 | logger.info(
195 | "===== Initialize %s with Xavier uniform distribution ====="
196 | % self.__class__.__name__
197 | )
198 | # NOTE: see https://github.com/pytorch/fairseq/blob/master/fairseq/modules/multihead_attention.py
199 | nn.init.xavier_uniform_(self.w_key.weight, gain=1 / math.sqrt(2))
200 | nn.init.xavier_uniform_(self.w_query.weight, gain=1 / math.sqrt(2))
201 | if bias:
202 | nn.init.constant_(self.w_key.bias, 0.0)
203 | nn.init.constant_(self.w_query.bias, 0.0)
204 |
205 | def reset(self):
206 | self.key = None
207 | self.mask = None
208 |
209 | def forward(
210 | self,
211 | key,
212 | query,
213 | mask,
214 | cache=False,
215 | boundary_leftmost=0,
216 | boundary_rightmost=10e6,
217 | ):
218 | """Compute chunkwise energy.
219 | Args:
220 | key (FloatTensor): `[B, klen, kdim]`
221 | query (FloatTensor): `[B, qlen, qdim]`
222 | mask (ByteTensor): `[B, qlen, klen]`
223 | cache (bool): cache key and mask
224 | Returns:
225 | e (FloatTensor): `[B, H_ca, qlen, klen]`
226 | """
227 | bs, klen, kdim = key.size()
228 | qlen = query.size(1)
229 |
230 | # Pre-computation of encoder-side features for computing scores
231 | if self.key is None or not cache:
232 | key = self.w_key(key).view(bs, -1, self.n_heads, self.d_k)
233 | self.key = key.transpose(2, 1).contiguous() # `[B, H_ca, klen, d_k]`
234 | self.mask = mask
235 | if mask is not None:
236 | self.mask = self.mask.unsqueeze(1).repeat(
237 | [1, self.n_heads, 1, 1]
238 | ) # `[B, H_ca, qlen, klen]`
239 | assert self.mask.size() == (bs, self.n_heads, qlen, klen), (
240 | self.mask.size(),
241 | (bs, self.n_heads, qlen, klen),
242 | )
243 |
244 | query = self.w_query(query).view(bs, -1, self.n_heads, self.d_k)
245 | query = query.transpose(2, 1).contiguous() # `[B, H_ca, qlen, d_k]`
246 | m = self.mask
247 |
248 | if self.atype == "add":
249 | k = self.key.unsqueeze(2) # `[B, H_ca, 1, klen, d_k]`
250 | # Truncate
251 | k = k[:, :, :, boundary_leftmost:boundary_rightmost]
252 | klen = k.size(3)
253 | if m is not None:
254 | m = m[:, :, :, boundary_leftmost:boundary_rightmost]
255 |
256 | r = torch.relu(k + query.unsqueeze(3)) # `[B, H_ca, qlen, klen, d_k]`
257 | r = (
258 | r.permute(0, 2, 3, 1, 4).contiguous().view(bs, qlen, klen, -1)
259 | ) # `[B, qlen, klen, H_ca * d_k]`
260 | r = self.v(r).permute(0, 3, 1, 2).contiguous() # `[B, H_ca, qlen, klen]`
261 | elif self.atype == "scaled_dot":
262 | k = self.key.transpose(3, 2)
263 | r = torch.matmul(query, k) / self.scale
264 |
265 | if m is not None:
266 | NEG_INF = float(np.finfo(torch.tensor(0, dtype=r.dtype).numpy().dtype).min)
267 | r = r.masked_fill_(m == 0, NEG_INF)
268 | assert r.size() == (bs, self.n_heads, qlen, klen), (
269 | r.size(),
270 | (bs, self.n_heads, qlen, klen),
271 | )
272 | return r
273 |
274 |
275 | class MonotonicAttention(nn.Module):
276 | def __init__(self, embed_dim, n_heads, dropout=0.1):
277 | super(MonotonicAttention, self).__init__()
278 |
279 | self.n_heads = n_heads
280 |
281 | self.monotonic_energy = MonotonicEnergy(
282 | kdim=embed_dim,
283 | qdim=embed_dim,
284 | adim=embed_dim,
285 | atype="scaled_dot",
286 | n_heads=n_heads,
287 | init_r=0,
288 | param_init="xavier_uniform",
289 | )
290 |
291 | self.chunk_energy = ChunkEnergy(
292 | kdim=embed_dim,
293 | qdim=embed_dim,
294 | adim=embed_dim,
295 | atype="scaled_dot",
296 | n_heads=n_heads,
297 | param_init="xavier_uniform",
298 | )
299 | self.dropout_attn = nn.Dropout(p=dropout) # for beta
300 | self.bd_offset = 0
301 |
302 | def forward(self, key, value, query, aw_prev=None, mask=None, cache=None):
303 |
304 | bs, klen = key.size()[:2]
305 | qlen = query.size(1)
306 |
307 | if aw_prev is None:
308 | # aw_prev = [1, 0, 0 ... 0]
309 | aw_prev = key.new_zeros(bs, self.n_heads_ma, 1, klen)
310 | aw_prev[:, :, :, 0:1] = key.new_ones(bs, self.n_heads_ma, 1, 1)
311 |
312 | # Compute monotonic energy
313 | e_ma = self.monotonic_energy(
314 | key, query, mask, cache=cache, boundary_leftmost=self.bd_offset
315 | ) # `[B, H_ma, qlen, klen]`
316 |
317 | p_choose = torch.sigmoid(
318 | add_gaussian_noise(e_ma, self.noise_std)
319 | ) # `[B, H_ma, qlen, klen]`
320 |
321 | # safe_cumprod computes cumprod in logspace with numeric checks
322 | cumprod_1mp_choose = safe_cumprod(
323 | 1 - p_choose, eps=self.eps
324 | ) # `[B, H_ma, qlen, klen]`
325 |
326 | # Compute recurrence relation solution
327 | alpha = []
328 | for i in range(qlen):
329 | denom = (
330 | 1
331 | if self.no_denom
332 | else torch.clamp(
333 | cumprod_1mp_choose[:, :, i : i + 1], min=self.eps, max=1.0
334 | )
335 | )
336 | aw_prev = (
337 | p_choose[:, :, i : i + 1]
338 | * cumprod_1mp_choose[:, :, i : i + 1]
339 | * torch.cumsum(aw_prev / denom, dim=-1)
340 | ) # `[B, H_ma, 1, klen]`
341 | # Mask the right part from the trigger point
342 | alpha.append(aw_prev)
343 |
344 | alpha = (
345 | torch.cat(alpha, dim=2) if qlen > 1 else alpha[-1]
346 | ) # `[B, H_ma, qlen, klen]`
347 | alpha_masked = alpha.clone()
348 |
349 | # Compute chunk energy
350 | e_ca = self.chunk_energy(
351 | key,
352 | query,
353 | mask,
354 | ) # `[B, (H_ma*)H_ca, qlen, ken]`
355 |
356 | beta = efficient_chunkwise_attention(
357 | alpha_masked,
358 | e_ca,
359 | mask,
360 | chunk_size=-1,
361 | n_heads_chunk=self.n_heads,
362 | sharpening_factor=1.0,
363 | share_chunkwise_attention=True,
364 | )
365 | beta = self.dropout_attn(beta)
366 |
367 | value = self.w_value(value).view(
368 | bs, -1, self.n_heads_ma * self.n_heads_ca, self.d_k
369 | )
370 | value = value.transpose(2, 1).contiguous() # `[B, H_ma * H_ca, klen, d_k]`
371 | cv = torch.matmul(
372 | alpha if self.w == 1 else beta, value
373 | ) # `[B, H_ma * H_ca, qlen, d_k]`
374 | cv = (
375 | cv.transpose(2, 1)
376 | .contiguous()
377 | .view(bs, -1, self.n_heads_ma * self.n_heads_ca * self.d_k)
378 | )
379 | cv = self.w_out(cv) # `[B, qlen, adim]`
380 |
381 |
382 | def add_gaussian_noise(xs, std):
383 | """Add Gaussian nosie to encourage discreteness."""
384 | noise = xs.new_zeros(xs.size()).normal_(std=std)
385 | return xs + noise
386 |
387 |
388 | def safe_cumprod(x, eps):
389 | """Numerically stable cumulative product by cumulative sum in log-space.
390 | Args:
391 | x (FloatTensor): `[B, H, qlen, klen]`
392 | Returns:
393 | x (FloatTensor): `[B, H, qlen, klen]`
394 | """
395 | return torch.exp(exclusive_cumsum(torch.log(torch.clamp(x, min=eps, max=1.0))))
396 |
397 |
398 | def exclusive_cumsum(x):
399 | """Exclusive cumulative summation [a, b, c] => [0, a, a + b].
400 | Args:
401 | x (FloatTensor): `[B, H, qlen, klen]`
402 | Returns:
403 | x (FloatTensor): `[B, H, qlen, klen]`
404 | """
405 | return torch.cumsum(
406 | torch.cat(
407 | [x.new_zeros(x.size(0), x.size(1), x.size(2), 1), x[:, :, :, :-1]], dim=-1
408 | ),
409 | dim=-1,
410 | )
411 |
412 |
413 | def exclusive_cumprod(x):
414 | """Exclusive cumulative product [a, b, c] => [1, a, a * b].
415 | Args:
416 | x (FloatTensor): `[B, H, qlen, klen]`
417 | Returns:
418 | x (FloatTensor): `[B, H, qlen, klen]`
419 | """
420 | return torch.cumprod(
421 | torch.cat(
422 | [x.new_ones(x.size(0), x.size(1), x.size(2), 1), x[:, :, :, :-1]], dim=-1
423 | ),
424 | dim=-1,
425 | )
426 |
427 |
428 | def efficient_chunkwise_attention(
429 | alpha,
430 | u,
431 | mask,
432 | chunk_size,
433 | n_heads_chunk,
434 | sharpening_factor,
435 | share_chunkwise_attention,
436 | ):
437 | """Compute chunkwise attention efficiently by clipping logits at training time.
438 | Args:
439 | alpha (FloatTensor): `[B, H_ma, qlen, klen]`
440 | u (FloatTensor): `[B, (H_ma*)H_ca, qlen, klen]`
441 | mask (ByteTensor): `[B, qlen, klen]`
442 | chunk_size (int): window size for chunkwise attention
443 | n_heads_chunk (int): number of chunkwise attention heads
444 | sharpening_factor (float): sharping factor for beta calculation
445 | share_chunkwise_attention (int): share CA heads among MA heads
446 | Returns:
447 | beta (FloatTensor): `[B, H_ma * H_ca, qlen, klen]`
448 | """
449 | bs, n_heads_mono, qlen, klen = alpha.size()
450 | alpha = alpha.unsqueeze(2) # `[B, H_ma, 1, qlen, klen]`
451 | u = u.unsqueeze(1) # `[B, 1, (H_ma*)H_ca, qlen, klen]`
452 | if n_heads_chunk > 1:
453 | alpha = alpha.repeat([1, 1, n_heads_chunk, 1, 1])
454 | if n_heads_mono > 1 and not share_chunkwise_attention:
455 | u = u.view(bs, n_heads_mono, n_heads_chunk, qlen, klen)
456 | # Shift logits to avoid overflow
457 | u -= torch.max(u, dim=-1, keepdim=True)[0]
458 | # Limit the range for numerical stability
459 | softmax_exp = torch.clamp(torch.exp(u), min=1e-5)
460 | # Compute chunkwise softmax denominators
461 | if chunk_size == -1:
462 | # infinite lookback attention
463 | # inner_items = alpha * sharpening_factor / torch.cumsum(softmax_exp, dim=-1)
464 | # beta = softmax_exp * torch.cumsum(inner_items.flip(dims=[-1]), dim=-1).flip(dims=[-1])
465 | # beta = beta.masked_fill(mask.unsqueeze(1), 0)
466 | # beta = beta / beta.sum(dim=-1, keepdim=True)
467 |
468 | softmax_denominators = torch.cumsum(softmax_exp, dim=-1)
469 | # Compute \beta_{i, :}. emit_probs are \alpha_{i, :}.
470 | beta = softmax_exp * moving_sum(
471 | alpha * sharpening_factor / softmax_denominators, back=0, forward=klen - 1
472 | )
473 | else:
474 | softmax_denominators = moving_sum(softmax_exp, back=chunk_size - 1, forward=0)
475 | # Compute \beta_{i, :}. emit_probs are \alpha_{i, :}.
476 | beta = softmax_exp * moving_sum(
477 | alpha * sharpening_factor / softmax_denominators,
478 | back=0,
479 | forward=chunk_size - 1,
480 | )
481 | return beta.view(bs, -1, qlen, klen)
482 |
483 |
484 | def moving_sum(x, back, forward):
485 | """Compute the moving sum of x over a chunk_size with the provided bounds.
486 | Args:
487 | x (FloatTensor): `[B, H_ma, H_ca, qlen, klen]`
488 | back (int):
489 | forward (int):
490 | Returns:
491 | x_sum (FloatTensor): `[B, H_ma, H_ca, qlen, klen]`
492 | """
493 | bs, n_heads_mono, n_heads_chunk, qlen, klen = x.size()
494 | x = x.view(-1, klen)
495 | # Moving sum is computed as a carefully-padded 1D convolution with ones
496 | x_padded = F.pad(
497 | x, pad=[back, forward]
498 | ) # `[B * H_ma * H_ca * qlen, back + klen + forward]`
499 | # Add a "channel" dimension
500 | x_padded = x_padded.unsqueeze(1)
501 | # Construct filters
502 | filters = x.new_ones(1, 1, back + forward + 1)
503 | x_sum = F.conv1d(x_padded, filters)
504 | x_sum = x_sum.squeeze(1).view(bs, n_heads_mono, n_heads_chunk, qlen, -1)
505 | return x_sum
506 |
--------------------------------------------------------------------------------
/src/model/positional_encoding.py:
--------------------------------------------------------------------------------
1 | import math
2 | import numpy as np
3 | import torch
4 | import torch.nn as nn
5 |
6 |
7 | class PositionalEncoding_2(nn.Module):
8 | r"""Inject some information about the relative or absolute position of the tokens
9 | in the sequence. The positional encodings have the same dimension as
10 | the embeddings, so that the two can be summed. Here, we use sine and cosine
11 | functions of different frequencies.
12 | .. math::
13 | \text{PosEncoder}(pos, 2i) = sin(pos/10000^(2i/d_model))
14 | \text{PosEncoder}(pos, 2i+1) = cos(pos/10000^(2i/d_model))
15 | \text{where pos is the word position and i is the embed idx)
16 | Args:
17 | d_model: the embed dim (required).
18 | dropout: the dropout value (default=0.1).
19 | max_len: the max. length of the incoming sequence (default=5000).
20 | """
21 |
22 | def __init__(self, d_model, dropout=0.1, max_len=5000):
23 | super(PositionalEncoding, self).__init__()
24 | self.dropout = nn.Dropout(p=dropout)
25 |
26 | pe = torch.zeros(max_len, d_model)
27 | position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
28 | div_term = torch.exp(
29 | torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
30 | )
31 | pe[:, 0::2] = torch.sin(position * div_term)
32 | pe[:, 1::2] = torch.cos(position * div_term)
33 | pe = pe.unsqueeze(0).transpose(0, 1)
34 | self.register_buffer("pe", pe)
35 |
36 | def forward(self, x):
37 | r"""Inputs of forward function
38 | Args:
39 | x: the sequence fed to the positional encoder model (required).
40 | Shape:
41 | x: [sequence length, batch size, embed dim]
42 | output: [sequence length, batch size, embed dim]
43 | Examples:
44 | """
45 |
46 | x = x + self.pe[: x.size(0), :]
47 | return self.dropout(x)
48 |
49 |
50 | class PositionalEncoding(nn.Module):
51 | r"""Inject some information about the relative or absolute position of the tokens
52 | in the sequence. The positional encodings have the same dimension as
53 | the embeddings, so that the two can be summed. Here, we use sine and cosine
54 | functions of different frequencies.
55 | .. math::
56 | \text{PosEncoder}(pos, 2i) = sin(pos/10000^(2i/d_model))
57 | \text{PosEncoder}(pos, 2i+1) = cos(pos/10000^(2i/d_model))
58 | \text{where pos is the word position and i is the embed idx)
59 | Args:
60 | d_model: the embed dim (required).
61 | dropout: the dropout value (default=0.1).
62 | max_len: the max. length of the incoming sequence (default=5000).
63 | """
64 |
65 | def __init__(self, d_model, dropout=0.1, max_len=5000):
66 | super(PositionalEncoding, self).__init__()
67 | self.dropout = nn.Dropout(p=dropout)
68 |
69 | pe = torch.zeros(max_len, d_model)
70 | position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
71 | div_term = torch.exp(
72 | torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
73 | )
74 | pe[:, 0::2] = torch.sin(position * div_term)
75 | pe[:, 1::2] = torch.sin(position * div_term)
76 | pe = pe.unsqueeze(0).transpose(0, 1)
77 | self.register_buffer("pe", pe)
78 |
79 | def forward(self, x):
80 | r"""Inputs of forward function
81 | Args:
82 | x: the sequence fed to the positional encoder model (required).
83 | Shape:
84 | x: [sequence length, batch size, embed dim]
85 | output: [sequence length, batch size, embed dim]
86 | Examples:
87 | """
88 |
89 | x = x + self.pe[: x.size(0), :]
90 | return self.dropout(x)
91 |
92 |
93 | class PositionalEncoding_00(nn.Module):
94 | r"""Inject some information about the relative or absolute position of the tokens
95 | in the sequence. The positional encodings have the same dimension as
96 | the embeddings, so that the two can be summed. Here, we use sine and cosine
97 | functions of different frequencies.
98 | .. math::
99 | \text{PosEncoder}(pos, 2i) = sin(pos/10000^(2i/d_model))
100 | \text{PosEncoder}(pos, 2i+1) = cos(pos/10000^(2i/d_model))
101 | \text{where pos is the word position and i is the embed idx)
102 | Args:
103 | d_model: the embed dim (required).
104 | dropout: the dropout value (default=0.1).
105 | max_len: the max. length of the incoming sequence (default=5000).
106 | """
107 |
108 | def __init__(self, d_model, max_len=5000):
109 | super(PositionalEncoding_00, self).__init__()
110 |
111 | pe = torch.zeros(max_len, d_model)
112 | position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
113 | div_term = torch.exp(
114 | torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
115 | )
116 | pe[:, 0::2] = torch.sin(position * div_term)
117 | pe[:, 1::2] = torch.cos(position * div_term)
118 | pe = pe.unsqueeze(0).transpose(0, 1)
119 | self.register_buffer("pe", pe)
120 |
121 | def forward(self, stop, step):
122 | r"""Inputs of forward function
123 | Args:
124 | x: the sequence fed to the positional encoder model (required).
125 | Shape:
126 | x: [sequence length, batch size, embed dim]
127 | output: [sequence length, batch size, embed dim]
128 | Examples:
129 | """
130 |
131 | return self.pe[0:stop:step, :]
132 |
133 |
134 | class PositionEmbeddingSine(nn.Module):
135 | """
136 | This is a more standard version of the position embedding, very similar to the one
137 | used by the Attention is all you need paper, generalized to work on images.
138 | """
139 |
140 | def __init__(
141 | self, num_pos_feats=64, temperature=10000, normalize=False, scale=None
142 | ):
143 | super().__init__()
144 | self.num_pos_feats = num_pos_feats
145 | self.temperature = temperature
146 | self.normalize = normalize
147 | if scale is not None and normalize is False:
148 | raise ValueError("normalize should be True if scale is passed")
149 | if scale is None:
150 | scale = 2 * math.pi
151 | self.scale = scale
152 |
153 | def forward(self, x):
154 | """
155 | x: [B, C, h, w]
156 | """
157 | B, C, h, w = x.shape
158 | ones = torch.ones_like(x[:, 0])
159 | y_embed = ones.cumsum(dim=1)
160 | x_embed = ones.cumsum(dim=2)
161 |
162 | if self.normalize:
163 | eps = 1e-6
164 | y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale
165 | x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale
166 |
167 | i = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device)
168 | dim_t = self.temperature ** (2 * (i // 2) / self.num_pos_feats)
169 |
170 | pos_x = x_embed[:, :, :, None] / dim_t
171 | pos_y = y_embed[:, :, :, None] / dim_t
172 | pos_x = torch.stack(
173 | [pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()], dim=4
174 | ).flatten(3)
175 | pos_y = torch.stack(
176 | (pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4
177 | ).flatten(3)
178 | pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)
179 | return pos
180 |
181 |
182 | class PositionalEncoding2DwDropout(nn.Module):
183 | def __init__(self, d_model, dropout=0.1, max_len=1000):
184 | super().__init__()
185 |
186 | self.dropout = nn.Dropout(p=dropout)
187 | self.d_model = d_model
188 |
189 | def forward(self, x, height, width):
190 | pe = torch.zeros(self.d_model, height, width, device=x.device)
191 |
192 | # Each dimension use half of d_model
193 | d_model = int(self.d_model / 2)
194 | div_term = torch.exp(
195 | torch.arange(0.0, d_model, 2) * -(math.log(10000.0) / d_model)
196 | ).cuda()
197 | pos_w = torch.arange(0.0, width // 2, device=x.device)
198 | pos_w = torch.cat([pos_w, pos_w.flip(dims=[0])]).unsqueeze(1)
199 | pos_h = torch.arange(0.0, height, device=x.device).unsqueeze(1)
200 | pe[0:d_model:2, :, :] = (
201 | torch.sin(pos_w * div_term)
202 | .transpose(0, 1)
203 | .unsqueeze(1)
204 | .repeat(1, height, 1)
205 | )
206 | pe[1:d_model:2, :, :] = (
207 | torch.cos(pos_w * div_term)
208 | .transpose(0, 1)
209 | .unsqueeze(1)
210 | .repeat(1, height, 1)
211 | )
212 |
213 | pe[d_model::2, :, :] = (
214 | torch.sin(pos_h * div_term).transpose(0, 1).unsqueeze(2).repeat(1, 1, width)
215 | )
216 | pe[d_model + 1 :: 2, :, :] = (
217 | torch.cos(pos_h * div_term).transpose(0, 1).unsqueeze(2).repeat(1, 1, width)
218 | )
219 |
220 | x = x + pe[None, :, :, :]
221 | return self.dropout(x)
222 |
223 |
224 | class PositionalEncoding3D(nn.Module):
225 | def __init__(self, channels):
226 | """
227 | :param channels: The last dimension of the tensor you want to apply pos emb to.
228 | """
229 | super(PositionalEncoding3D, self).__init__()
230 | channels = int(np.ceil(channels / 3))
231 | if channels % 2:
232 | channels += 1
233 | self.channels = channels
234 | inv_freq = 1.0 / (10000 ** (torch.arange(0, channels, 2).float() / channels))
235 | self.register_buffer("inv_freq", inv_freq)
236 |
237 | def forward(self, tensor):
238 | """
239 | :param tensor: A 5d tensor of size (batch_size, x, y, z, ch)
240 | :return: Positional Encoding Matrix of size (batch_size, x, y, z, ch)
241 | """
242 | if len(tensor.shape) != 5:
243 | raise RuntimeError("The input tensor has to be 5d!")
244 |
245 | _, x, y, z, orig_ch = tensor.shape
246 |
247 | pos_x = torch.arange(x, device=tensor.device).type(self.inv_freq.type())
248 | pos_y = torch.arange(y, device=tensor.device).type(self.inv_freq.type())
249 | pos_z = torch.arange(z, device=tensor.device).type(self.inv_freq.type())
250 | sin_inp_x = torch.einsum("i,j->ij", pos_x, self.inv_freq)
251 | sin_inp_y = torch.einsum("i,j->ij", pos_y, self.inv_freq)
252 | sin_inp_z = torch.einsum("i,j->ij", pos_z, self.inv_freq)
253 | emb_x = (
254 | torch.cat((sin_inp_x.sin(), sin_inp_x.cos()), dim=-1)
255 | .unsqueeze(1)
256 | .unsqueeze(1)
257 | )
258 | emb_y = torch.cat((sin_inp_y.sin(), sin_inp_y.cos()), dim=-1).unsqueeze(1)
259 | emb_z = torch.cat((sin_inp_z.sin(), sin_inp_z.cos()), dim=-1)
260 | emb = torch.zeros((x, y, z, self.channels * 3), device=tensor.device).type(
261 | tensor.type()
262 | )
263 | emb[:, :, :, : self.channels] = emb_x
264 | emb[:, :, :, self.channels : 2 * self.channels] = emb_y
265 | emb[:, :, :, 2 * self.channels :] = emb_z
266 |
267 | return emb[None, :, :, :, :orig_ch]
268 |
--------------------------------------------------------------------------------
9 |