Commit 907ca927 authored by Myle Ott's avatar Myle Ott
Browse files

Better support for torch.no_grad (since volatile is deprecated)

parent 0b84ab19

fairseq/modules/linearized_convolution.py

+1 −0
Changes for fairseq/modules/linearized_convolution.py: 1 added line, 0 removed lines.
Original line number Diff line number Diff line
@@ -69,6 +69,7 @@ class LinearizedConvolution(ConvTBC): 
            # append next input

            self.input_buffer[:, -1, :] = input[:, -1, :]

            input = utils.volatile_variable(self.input_buffer)

        with utils.maybe_no_grad():

            output = F.linear(input.view(bsz, -1), weight, self.bias)

        return output.view(bsz, 1, -1)

fairseq/sequence_generator.py

+2 −0
Changes for fairseq/sequence_generator.py: 2 added lines, 0 removed lines.
Original line number Diff line number Diff line
@@ -71,6 +71,7 @@ class SequenceGenerator(object): 
            srclen = input['src_tokens'].size(1)

            if timer is not None:

                timer.start()

            with utils.maybe_no_grad():

                hypos = self.generate(input['src_tokens'], beam_size=beam_size,

                                      maxlen=int(maxlen_a*srclen + maxlen_b))

            if timer is not None:
@@ -327,6 +328,7 @@ class SequenceGenerator(object): 
        avg_probs = None

        avg_attn = None

        for model, encoder_out in zip(self.models, encoder_outs):

            with utils.maybe_no_grad():

                decoder_out, attn = model.decoder(tokens, encoder_out)

            probs = model.get_normalized_probs(decoder_out[:, -1, :], log_probs=False).data

            if avg_probs is None:

fairseq/utils.py

+3 −2
Changes for fairseq/utils.py: 3 added lines, 2 removed lines.
Original line number Diff line number Diff line
@@ -176,7 +176,7 @@ def _upgrade_args(args): 
    return args





def maybe_no_grad(condition):

def maybe_no_grad(condition=True):

    if hasattr(torch, 'no_grad') and condition:

        return torch.no_grad()

    # no-op context manager
@@ -185,8 +185,9 @@ def maybe_no_grad(condition): 


def volatile_variable(*args, **kwargs):

    if hasattr(torch, 'no_grad'):

        with torch.no_grad():

        # volatile has been deprecated, use the no_grad context manager instead

        return Variable(*args, **kwargs)

    else:

        return Variable(*args, **kwargs, volatile=True)