jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def forward(self, x):
        """Forward pass of the network.
        """
        N = x.size()[0]
        x1 = self.backbone.features(x)

        H = x1.size()[2]
        W = x1.size()[3]
        assert x1.size()[1] == 512

        x2 = self.sfeatures(x)
        H2 = x2.size()[2]
        W2 = x2.size()[3]
        assert x2.size()[1] == 128

        if (H != H2) | (W != W2):
            x2 = F.interpolate(x2, (H, W), mode='bilinear')

        x1 = x1.view(N, 512, H * W)
        x2 = x2.view(N, 128, H * W)
        x = torch.bmm(x1, torch.transpose(x2, 1, 2)) / (H * W)  # Bilinear
        assert x.size() == (N, 512, 128)
        x = x.view(N, 512 * 128)
        #x = torch.sqrt(x + 1e-8)
        x = F.normalize(x)

        output = []

        if not self.config.JL:
            for idx, fc in enumerate(self.fc):
                if not self.config.train:
                    for W in fc.parameters():
                        fc.weight.data = F.normalize(W, p=2, dim=1)
                output.append(fc(x))
        else:
            for idx, fc in enumerate(self.fc):
                output.append(fc(x))

        sfeat = 0
        return output, sfeat
//  [python]
def forward(self, *in_x):
        features = self.encoder(*in_x)
        b, c, h, w = in_x[0].shape
        features = features[1:]
        ret = []
        for i in range(len(features)):
            conv = getattr(self, 'shrink%d' % i)
            x = conv(features[i])
            ratio = h // features[i].shape[-2]
            ret.append(F.upsample_bilinear(x, scale_factor=ratio))
        ret = torch.cat(ret, 1)
        ret = self.aggregate(ret)
        ret = self.bn(ret)
        ret = F.relu(ret, True)
        logits = self.outc(ret)
        return logits
//  [python]
def forward(self, *in_x):
        features = self.encoder(*in_x)
        b, c, h, w = in_x[0].shape
        features = features[1:]
        ret = []
        for i in range(len(features)):
            conv = getattr(self, 'shrink%d' % i)
            x = conv(features[i])
            ratio = h // features[i].shape[-2]
            ret.append(F.interpolate(x, scale_factor=ratio, mode='bilinear'))
        ret = torch.cat(ret, 1)
        ret = self.aggregate(ret)
        ret = self.bn(ret)
        ret = F.relu(ret, True)
        logits = self.outc(ret)
        return logits
//  [python]
def forward(self, x):
        """Forward pass of the network.
        """
        N = x.size()[0]

        x1 = self.backbone.conv1(x)
        x1 = self.backbone.bn1(x1)
        x1 = self.backbone.relu(x1)
        x1 = self.backbone.maxpool(x1)

        x1 = self.backbone.layer1(x1)
        x1 = self.backbone.layer2(x1)
        x1 = self.backbone.layer3(x1)
        x1 = self.backbone.layer4(x1)

        H = x1.size()[2]
        W = x1.size()[3]
        assert x1.size()[1] == 512

        x2 = self.sfeatures(x)
        H2 = x2.size()[2]
        W2 = x2.size()[3]
        assert x2.size()[1] == 128

        sfeat = self.pooling(x2)
        sfeat = sfeat.squeeze(3).squeeze(2)
        sfeat = F.normalize(sfeat, p=2)

        if (H != H2) | (W != W2):
            x2 = F.upsample_bilinear(x2, (H, W))

        x1 = x1.view(N, 512, H * W)
        x2 = x2.view(N, 128, H * W)
        x = torch.bmm(x1, torch.transpose(x2, 1, 2)) / (H * W)  # Bilinear
        assert x.size() == (N, 512, 128)
        x = x.view(N, 512 * 128)
        #x = torch.sqrt(x + 1e-8)
        x = F.normalize(x)

        output = []

        if not self.config.JL:
            for idx, fc in enumerate(self.fc):
                if not self.config.train:
                    for W in fc.parameters():
                        fc.weight.data = F.normalize(W, p=2, dim=1)
                output.append(fc(x))
        else:
            for idx, fc in enumerate(self.fc):
                output.append(fc(x))


        return output, sfeat
//  [python]
def forward(self, x):
        """Forward pass of the network.
        """
        N = x.size()[0]

        x1 = self.backbone.conv1(x)
        x1 = self.backbone.bn1(x1)
        x1 = self.backbone.relu(x1)
        x1 = self.backbone.maxpool(x1)

        x1 = self.backbone.layer1(x1)
        x1 = self.backbone.layer2(x1)
        x1 = self.backbone.layer3(x1)
        x1 = self.backbone.layer4(x1)

        H = x1.size()[2]
        W = x1.size()[3]
        assert x1.size()[1] == 512

        x2 = self.sfeatures(x)
        H2 = x2.size()[2]
        W2 = x2.size()[3]
        assert x2.size()[1] == 128

        sfeat = self.pooling(x2)
        sfeat = sfeat.squeeze(3).squeeze(2)
        sfeat = F.normalize(sfeat, p=2)

        if (H != H2) | (W != W2):
            x2 = F.interpolate(x2, (H, W), mode='bilinear')

        x1 = x1.view(N, 512, H * W)
        x2 = x2.view(N, 128, H * W)
        x = torch.bmm(x1, torch.transpose(x2, 1, 2)) / (H * W)  # Bilinear
        assert x.size() == (N, 512, 128)
        x = x.view(N, 512 * 128)
        #x = torch.sqrt(x + 1e-8)
        x = F.normalize(x)

        output = []

        if not self.config.JL:
            for idx, fc in enumerate(self.fc):
                if not self.config.train:
                    for W in fc.parameters():
                        fc.weight.data = F.normalize(W, p=2, dim=1)
                output.append(fc(x))
        else:
            for idx, fc in enumerate(self.fc):
                output.append(fc(x))


        return output, sfeat
//  [python]
def forward(self, *in_x):
        features = self.encoder(*in_x)
        features = features[1:]
        x = features[-1]
        for i in range(len(features) - 2, -1, -1):
            conv = getattr(self, 'shrink%d' % i)
            x = F.upsample_nearest(x, scale_factor=2)
            x = conv(x)
            if features[i].shape[-1] != x.shape[-1]:
                x2 = F.upsample_nearest(features[i], scale_factor=2)
            else:
                x2 = features[i]
            x = torch.cat([x, x2], 1)
            conv2 = getattr(self, 'shrink2%d' % i)
            x = conv2(x)
            bn = getattr(self, 'batchnorm%d' % i)
            x = bn(x)
            x = F.relu(x)
            x = self.dropout(x)
        x = F.upsample_nearest(x, scale_factor=2)
        logits = self.outc(x)
        return logits
//  [python]
def forward(self, *in_x):
        features = self.encoder(*in_x)
        features = features[1:]
        x = features[-1]
        for i in range(len(features) - 2, -1, -1):
            conv = getattr(self, 'shrink%d' % i)
            x = F.interpolate(x, scale_factor=2, mode='nearest')
            x = conv(x)
            if features[i].shape[-1] != x.shape[-1]:
                x2 = F.interpolate(features[i], scale_factor=2, mode='nearest')
            else:
                x2 = features[i]
            x = torch.cat([x, x2], 1)
            conv2 = getattr(self, 'shrink2%d' % i)
            x = conv2(x)
            bn = getattr(self, 'batchnorm%d' % i)
            x = bn(x)
            x = F.relu(x)
            x = self.dropout(x)
        x = F.interpolate(x, scale_factor=2, mode='nearest')
        logits = self.outc(x)
        return logits
//  [python]
def forward(self, x):
        """
        Arguments:
            x (list[Tensor]): feature maps for each feature level.
        Returns:
            results (tuple[Tensor]): feature maps after FPN layers.
                They are ordered from highest resolution first.
        """
        last_inner = getattr(self, self.inner_blocks[-1])(x[-1])
        results = []
        results.append(getattr(self, self.layer_blocks[-1])(last_inner))
        for feature, inner_block, layer_block in zip(
            x[:-1][::-1], self.inner_blocks[:-1][::-1], self.layer_blocks[:-1][::-1]
        ):
            if not inner_block:
                continue
            # inner_top_down = F.interpolate(last_inner, scale_factor=2, mode="nearest")
            inner_lateral = getattr(self, inner_block)(feature)
            # TODO use size instead of scale to make it robust to different sizes
            inner_top_down = F.upsample(last_inner, size=inner_lateral.shape[-2:],
            mode='bilinear', align_corners=False)
            last_inner = inner_lateral + inner_top_down
            results.insert(0, getattr(self, layer_block)(last_inner))

        if isinstance(self.top_blocks, LastLevelP6P7):
            last_results = self.top_blocks(x[-1], results[-1])
            results.extend(last_results)
        elif isinstance(self.top_blocks, LastLevelMaxPool):
            last_results = self.top_blocks(results[-1])
            results.extend(last_results)

        return tuple(results)
//  [python]
def forward(self, x):
        """
        Arguments:
            x (list[Tensor]): feature maps for each feature level.
        Returns:
            results (tuple[Tensor]): feature maps after FPN layers.
                They are ordered from highest resolution first.
        """
        last_inner = getattr(self, self.inner_blocks[-1])(x[-1])
        results = []
        results.append(getattr(self, self.layer_blocks[-1])(last_inner))
        for feature, inner_block, layer_block in zip(
            x[:-1][::-1], self.inner_blocks[:-1][::-1], self.layer_blocks[:-1][::-1]
        ):
            if not inner_block:
                continue
            # inner_top_down = F.interpolate(last_inner, scale_factor=2, mode="nearest")
            inner_lateral = getattr(self, inner_block)(feature)
            # TODO use size instead of scale to make it robust to different sizes
            inner_top_down = F.interpolate(last_inner, size=inner_lateral.shape[-2:],
            mode='bilinear', align_corners=False)
            last_inner = inner_lateral + inner_top_down
            results.insert(0, getattr(self, layer_block)(last_inner))

        if isinstance(self.top_blocks, LastLevelP6P7):
            last_results = self.top_blocks(x[-1], results[-1])
            results.extend(last_results)
        elif isinstance(self.top_blocks, LastLevelMaxPool):
            last_results = self.top_blocks(results[-1])
            results.extend(last_results)

        return tuple(results)
//  [python]
def forward(self,
              input,
              size=None,
              scale_factor=None,
              mode='nearest',
              align_corners=None):
    return F.upsample(input, size, scale_factor, mode, align_corners)
//  [python]
def forward(self,
              input,
              size=None,
              scale_factor=None,
              mode='nearest',
              align_corners=None):
    return F.interpolate(input, size, scale_factor, mode, align_corners)
//  [python]
def forward(self, x):
        x = self.conv2d(x)
        x = F.upsample(x, scale_factor=2)
        x = self.up(x)
        return x
//  [python]
def forward(self, x):
        x = self.conv2d(x)
        x = F.interpolate(x, scale_factor=2)
        x = self.up(x)
        return x
//  [python]
def img_crop(input, size=224):
    input_cropped = F.upsample(input, size=(size, size), mode='bilinear', align_corners=True)
    return input_cropped
//  [python]
def img_crop(input, size=224):
    input_cropped = F.interpolate(input, size=(size, size), mode='bilinear', align_corners=True)
    return input_cropped
//  [python]
def compute_clip_loss(img, text):
    img = torch.nn.functional.upsample_bilinear(img, (224, 224))
    tokenized_text = clip.tokenize([text]).to(device)

    img_logits, _text_logits = clip_model(img, tokenized_text)

    return 1/img_logits * 100
//  [python]
def compute_clip_loss(img, text):
    img = torch.nn.functional.interpolate(img, (224, 224), mode='bilinear')
    tokenized_text = clip.tokenize([text]).to(device)

    img_logits, _text_logits = clip_model(img, tokenized_text)

    return 1/img_logits * 100
//  [python]
def forward(self, x):
        input = x
        x = self.conv1(x)
        x = self.bn0(x)
        x = self.relu(x)
        conv_x = x
        x = self.maxpool(x)
        pool_x = x

        fm1 = self.layer1(x)
        fm2 = self.layer2(fm1)
        fm3 = self.layer3(fm2)
        fm4 = self.layer4(fm3)

        gcfm1 = self.gcn1(fm4)
        gcfm2 = self.gcn2(fm3)
        gcfm3 = self.gcn3(fm2)
        gcfm4 = self.gcn4(pool_x)
        gcfm5 = self.gcn5(conv_x)

        fs1 = F.upsample_bilinear(gcfm1, fm3.size()[2:]) + gcfm2
        fs2 = F.upsample_bilinear(fs1, fm2.size()[2:]) + gcfm3
        fs3 = F.upsample_bilinear(fs2, pool_x.size()[2:]) + gcfm4
        fs4 = F.upsample_bilinear(fs3, conv_x.size()[2:]) + gcfm5
        out = F.upsample_bilinear(fs4, input.size()[2:])

        return out
//  [python]
def forward(self, x):
        input = x
        x = self.conv1(x)
        x = self.bn0(x)
        x = self.relu(x)
        conv_x = x
        x = self.maxpool(x)
        pool_x = x

        fm1 = self.layer1(x)
        fm2 = self.layer2(fm1)
        fm3 = self.layer3(fm2)
        fm4 = self.layer4(fm3)

        gcfm1 = self.gcn1(fm4)
        gcfm2 = self.gcn2(fm3)
        gcfm3 = self.gcn3(fm2)
        gcfm4 = self.gcn4(pool_x)
        gcfm5 = self.gcn5(conv_x)

        fs1 = F.interpolate(gcfm1, size=fm3.size()[2:], mode='bilinear') + gcfm2
        fs2 = F.interpolate(fs1, size=fm2.size()[2:], mode='bilinear') + gcfm3
        fs3 = F.interpolate(fs2, size=pool_x.size()[2:], mode='bilinear') + gcfm4
        fs4 = F.interpolate(fs3, size=conv_x.size()[2:], mode='bilinear') + gcfm5
        out = F.interpolate(fs4, size=input.size()[2:], mode='bilinear')

        return out
//  [python]
def bilinear_attention_pooling(self, features, attentions):
        B, C, H, W = features.size()
        _, M, AH, AW = attentions.size()

        # match size
        if AH != H or AW != W:
            attentions = F.upsample_bilinear(attentions, size=(H, W))

        # feature_matrix: (B, M, C) -> (B, M * C)
        if self.pooling is None:
            feature_matrix = (torch.einsum('imjk,injk->imn', (attentions, features)) / float(H * W)).view(B, -1)
        else:
            feature_matrix = []
            for i in range(M):
                AiF = self.pooling(features * attentions[:, i:i + 1, ...]).view(B, -1)
                feature_matrix.append(AiF)
            feature_matrix = torch.cat(feature_matrix, dim=1)

        # sign-sqrt
        feature_matrix_raw = torch.sign(feature_matrix) * torch.sqrt(torch.abs(feature_matrix) + self.EPSILON)

        # l2 normalization along dimension M and C
        feature_matrix = F.normalize(feature_matrix_raw, dim=-1)

        return feature_matrix
//  [python]
def bilinear_attention_pooling(self, features, attentions):
        B, C, H, W = features.size()
        _, M, AH, AW = attentions.size()

        # match size
        if AH != H or AW != W:
            attentions = F.interpolate(attentions, size=(H, W), mode='bilinear')

        # feature_matrix: (B, M, C) -> (B, M * C)
        if self.pooling is None:
            feature_matrix = (torch.einsum('imjk,injk->imn', (attentions, features)) / float(H * W)).view(B, -1)
        else:
            feature_matrix = []
            for i in range(M):
                AiF = self.pooling(features * attentions[:, i:i + 1, ...]).view(B, -1)
                feature_matrix.append(AiF)
            feature_matrix = torch.cat(feature_matrix, dim=1)

        # sign-sqrt
        feature_matrix_raw = torch.sign(feature_matrix) * torch.sqrt(torch.abs(feature_matrix) + self.EPSILON)

        # l2 normalization along dimension M and C
        feature_matrix = F.normalize(feature_matrix_raw, dim=-1)

        return feature_matrix
//  [python]
def masking(input, target=None, logit_t=None, logit_s=None, model_t=None, model_s=None, opt=None):
    r"""Similar interface to cutout.
    This implements the masking operation in MAE (https://arxiv.org/abs/2111.06377).
    """
    if opt.dataset in ['mnist', 'cifar10', 'cifar100', 'svhn']:
        img_size = 32
    elif opt.dataset in ['tinyimagenet']:
        img_size = 64
    elif opt.dataset in ['imagenet', 'imagenet100']:
        img_size = 224
    else:
        raise NotImplementedError
    n_patch = opt.n_patch 
    mask = torch.randperm(n_patch * n_patch) >= int(opt.mask_zero_ratio * n_patch * n_patch) # This presumes the input is square
    mask = mask.view(n_patch, n_patch) # [n_patch, n_patch]
    mask = mask.float().unsqueeze(dim=0).unsqueeze(dim=0) # [1, 1, n_patch, n_patch]
    mask = F.upsample_nearest(mask, scale_factor=img_size//n_patch).cuda()
    input_mix = input * mask
    target_mix, logit_t_mix, logit_s_mix, rand_index, lam = None, None, None, None, None # placeholder to maintain interface
    return input_mix, target_mix, logit_t_mix, logit_s_mix, rand_index, lam
//  [python]
def masking(input, target=None, logit_t=None, logit_s=None, model_t=None, model_s=None, opt=None):
    r"""Similar interface to cutout.
    This implements the masking operation in MAE (https://arxiv.org/abs/2111.06377).
    """
    if opt.dataset in ['mnist', 'cifar10', 'cifar100', 'svhn']:
        img_size = 32
    elif opt.dataset in ['tinyimagenet']:
        img_size = 64
    elif opt.dataset in ['imagenet', 'imagenet100']:
        img_size = 224
    else:
        raise NotImplementedError
    n_patch = opt.n_patch 
    mask = torch.randperm(n_patch * n_patch) >= int(opt.mask_zero_ratio * n_patch * n_patch) # This presumes the input is square
    mask = mask.view(n_patch, n_patch) # [n_patch, n_patch]
    mask = mask.float().unsqueeze(dim=0).unsqueeze(dim=0) # [1, 1, n_patch, n_patch]
    mask = F.interpolate(mask, scale_factor=img_size//n_patch, mode='nearest').cuda()
    input_mix = input * mask
    target_mix, logit_t_mix, logit_s_mix, rand_index, lam = None, None, None, None, None # placeholder to maintain interface
    return input_mix, target_mix, logit_t_mix, logit_s_mix, rand_index, lam
//  [python]
def _upsample(self, x, size, scale=1):
        _, _, H, W = size
        return F.upsample(x, size=(H // scale, W // scale), mode='bilinear')
//  [python]
def _upsample(self, x, size, scale=1):
        _, _, H, W = size
        return F.interpolate(x, size=(H // scale, W // scale), mode='bilinear')
//  [python]
def forward(self, x, phase=None):
        feat = []
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)
        x = self.layer1(x)              #(2L, 256L, 129L, 129L)
        x = self.layer2(x)              #(2L, 512L, 65L, 65L)
        feat += [x]
        feat0 = x
        x = self.layer3(x)              #(2L, 1024L, 33L, 33L)
        feat += [x]
        feat1 = x
        x = self.layer4(x)
        feat += [x]
        feat2 = x
        for k, v in enumerate(self.extra_layers):
            x = v(x)
            if k in [5, 11, 17, 23]:
                feat += [x]
        if self.att == True:
            ######### fusion #################################################
            feat0 = self.fusion_layers[0](feat0)
            feat1 = F.upsample_bilinear(self.fusion_layers[1](feat1),size=(65,65))
            feat2 = F.upsample_bilinear(self.fusion_layers[2](feat2),size=(65,65))
            feat[0] = F.relu(self.fusion_conv(self.fusion_bn(torch.cat([feat0, feat1, feat2], dim=1))))
            #################################################################
            feat_new = []
            for (x,l) in zip(feat, self.att_layers):
                feat_new.append(l(x))
        ########## PreEnd #################################################
        locs = []
        conf = []
        for (x, l, c) in zip(feat, self.locs_layers, self.conf_layers):
            locs += [l(x).permute(0,2,3,1).contiguous()]
            conf += [c(x).permute(0,2,3,1).contiguous()]

        locs = torch.cat([o.view(o.size(0),-1) for o in locs], dim=1)
        conf = torch.cat([o.view(o.size(0),-1) for o in conf], dim=1)

        if phase == 'test':
            output = self.detect(locs.view(locs.size(0), -1, 4),
                                 self.softmax(conf.view(-1, self.num_classes)),
                                 self.prior_boxes.type(type(x.data)))
        else:
            output = (locs.view(locs.size(0),-1,4),
                      conf.view(conf.size(0),-1,self.num_classes),
                      self.prior_boxes)
        return output
//  [python]
def forward(self, x, phase=None):
        feat = []
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)
        x = self.layer1(x)              #(2L, 256L, 129L, 129L)
        x = self.layer2(x)              #(2L, 512L, 65L, 65L)
        feat += [x]
        feat0 = x
        x = self.layer3(x)              #(2L, 1024L, 33L, 33L)
        feat += [x]
        feat1 = x
        x = self.layer4(x)
        feat += [x]
        feat2 = x
        for k, v in enumerate(self.extra_layers):
            x = v(x)
            if k in [5, 11, 17, 23]:
                feat += [x]
        if self.att == True:
            ######### fusion #################################################
            feat0 = self.fusion_layers[0](feat0)
            feat1 = F.interpolate(self.fusion_layers[1](feat1),size=(65,65), mode='bilinear', align_corners=False)
            feat2 = F.interpolate(self.fusion_layers[2](feat2),size=(65,65), mode='bilinear', align_corners=False)
            feat[0] = F.relu(self.fusion_conv(self.fusion_bn(torch.cat([feat0, feat1, feat2], dim=1))))
            #################################################################
            feat_new = []
            for (x,l) in zip(feat, self.att_layers):
                feat_new.append(l(x))
        ########## PreEnd #################################################
        locs = []
        conf = []
        for (x, l, c) in zip(feat, self.locs_layers, self.conf_layers):
            locs += [l(x).permute(0,2,3,1).contiguous()]
            conf += [c(x).permute(0,2,3,1).contiguous()]

        locs = torch.cat([o.view(o.size(0),-1) for o in locs], dim=1)
        conf = torch.cat([o.view(o.size(0),-1) for o in conf], dim=1)

        if phase == 'test':
            output = self.detect(locs.view(locs.size(0), -1, 4),
                                 self.softmax(conf.view(-1, self.num_classes)),
                                 self.prior_boxes.type(type(x.data)))
        else:
            output = (locs.view(locs.size(0),-1,4),
                      conf.view(conf.size(0),-1,self.num_classes),
                      self.prior_boxes)
        return output
//  [python]
def forward(self, x):
        next_x = {}
        feature_names = list(x.keys())
        for level, name in enumerate(feature_names):

            feature = x[name]

            offset_mask = self.offset(feature)
            offset = offset_mask[:, :18, :, :]
            mask = offset_mask[:, 18:, :, :].sigmoid()
            conv_args = dict(offset=offset, mask=mask)

            temp_fea = [self.DyConv[1](feature, **conv_args)]
            if level > 0:
                temp_fea.append(self.DyConv[2](x[feature_names[level - 1]], **conv_args))
            if level < len(x) - 1:
                temp_fea.append(F.upsample_bilinear(self.DyConv[0](x[feature_names[level + 1]], **conv_args),
                                                    size=[feature.size(2), feature.size(3)]))
            attn_fea = []
            res_fea = []
            for fea in temp_fea:
                res_fea.append(fea)
                attn_fea.append(self.AttnConv(fea))

            res_fea = torch.stack(res_fea)
            spa_pyr_attn = self.h_sigmoid(torch.stack(attn_fea))
            mean_fea = torch.mean(res_fea * spa_pyr_attn, dim=0, keepdim=False)
            next_x[name] = self.relu(mean_fea)

        return next_x
//  [python]
def forward(self, x):
        next_x = {}
        feature_names = list(x.keys())
        for level, name in enumerate(feature_names):

            feature = x[name]

            offset_mask = self.offset(feature)
            offset = offset_mask[:, :18, :, :]
            mask = offset_mask[:, 18:, :, :].sigmoid()
            conv_args = dict(offset=offset, mask=mask)

            temp_fea = [self.DyConv[1](feature, **conv_args)]
            if level > 0:
                temp_fea.append(self.DyConv[2](x[feature_names[level - 1]], **conv_args))
            if level < len(x) - 1:
                temp_fea.append(F.interpolate(self.DyConv[0](x[feature_names[level + 1]], **conv_args),
                                                    size=[feature.size(2), feature.size(3)], mode='bilinear'))
            attn_fea = []
            res_fea = []
            for fea in temp_fea:
                res_fea.append(fea)
                attn_fea.append(self.AttnConv(fea))

            res_fea = torch.stack(res_fea)
            spa_pyr_attn = self.h_sigmoid(torch.stack(attn_fea))
            mean_fea = torch.mean(res_fea * spa_pyr_attn, dim=0, keepdim=False)
            next_x[name] = self.relu(mean_fea)

        return next_x
//  [python]
def ARM(self,c1,c2,c3):
        _,_,s3,s4=c2.size()
        c2=self.ARM_layer1(F.avg_pool2d(c2,kernel_size=(s3,s4)))*c2+F.upsample_bilinear(c1,scale_factor=2)
        _, _, s3, s4 = c3.size()
        c3 = self.ARM_layer2(F.avg_pool2d(c3,kernel_size=(s3,s4))) * c3 + F.upsample_bilinear(c2,scale_factor=2)
        return c3
//  [python]
def ARM(self,c1,c2,c3):
        _,_,s3,s4=c2.size()
        c2=self.ARM_layer1(F.avg_pool2d(c2,kernel_size=(s3,s4)))*c2+F.interpolate(c1,scale_factor=2,mode='bilinear')
        _, _, s3, s4 = c3.size()
        c3 = self.ARM_layer2(F.avg_pool2d(c3,kernel_size=(s3,s4))) * c3 + F.interpolate(c2,scale_factor=2,mode='bilinear')
        return c3
//  [python]
def forward(self, x):
        context_input=x

        c1,c2,c3=self.Context_model(context_input)

        c1=self.enc1(c1)
        c2=self.enc2(c2)
        c_feat=self.ARM(c1,c2,c3)
        x=self.Spatial_entry(x)

        for layer in self.Spatial_model:
            x=layer(x)

        s_out=self.enc3(x)

        feat=self.FFM(c_feat,s_out)

        score=self.classifer(feat)

        score=F.upsample_bilinear(score,scale_factor=8)

        return score
//  [python]
def forward(self, x):
        context_input=x

        c1,c2,c3=self.Context_model(context_input)

        c1=self.enc1(c1)
        c2=self.enc2(c2)
        c_feat=self.ARM(c1,c2,c3)
        x=self.Spatial_entry(x)

        for layer in self.Spatial_model:
            x=layer(x)

        s_out=self.enc3(x)

        feat=self.FFM(c_feat,s_out)

        score=self.classifer(feat)

        score=F.interpolate(score,scale_factor=8,mode='bilinear')

        return score
//  [python]
def forward(self, x):
        x4, x3, x2, x1, x0 = self.model(x)

        a = self.enc4_1(x4)
        b = self.enc3_1(x3)
        c = self.enc2_1(x2)
        d = self.enc1_1(x1)

        b = self.sigmoid(F.upsample_bilinear(a, scale_factor=2)) * self.enc_1(b) + F.upsample_bilinear(a,
                                                                                                       scale_factor=2)
        c = self.sigmoid(F.upsample_bilinear(b, scale_factor=2)) * self.enc_2(c) + F.upsample_bilinear(b,
                                                                                                       scale_factor=2)
        d = self.sigmoid(F.upsample_bilinear(c, scale_factor=2)) * self.enc_3(d) + F.upsample_bilinear(c,
                                                                                                       scale_factor=2)

        d = self.relu(self.res1(d) + d)

        d = self.bn1(d * (1 + F.tanh(torch.matmul(d, d.transpose(3, 2)) @ d)))

        score = self.classifer(d)
        score = F.upsample_bilinear(score, scale_factor=4)

        return score
//  [python]
def forward(self, x):
        x4, x3, x2, x1, x0 = self.model(x)

        a = self.enc4_1(x4)
        b = self.enc3_1(x3)
        c = self.enc2_1(x2)
        d = self.enc1_1(x1)

        b = self.sigmoid(F.interpolate(a, scale_factor=2, mode='bilinear')) * self.enc_1(b) + F.interpolate(a, scale_factor=2, mode='bilinear')
        c = self.sigmoid(F.interpolate(b, scale_factor=2, mode='bilinear')) * self.enc_2(c) + F.interpolate(b, scale_factor=2, mode='bilinear')
        d = self.sigmoid(F.interpolate(c, scale_factor=2, mode='bilinear')) * self.enc_3(d) + F.interpolate(c, scale_factor=2, mode='bilinear')

        d = self.relu(self.res1(d) + d)

        d = self.bn1(d * (1 + F.tanh(torch.matmul(d, d.transpose(3, 2)) @ d)))

        score = self.classifer(d)
        score = F.interpolate(score, scale_factor=4, mode='bilinear')

        return score
//  [python]
def forward(self,x):
     
        x4,x3,x2,x1,x0=self.model(x)

        a=self.enc4_1(x4)
        b=self.enc3_1(x3)
        c=self.enc2_1(x2)
        d=self.enc1_1(x1)

        b=self.sigmoid(F.upsample_bilinear(a,scale_factor=2))*self.enc_1(b)+F.upsample_bilinear(a,scale_factor=2)
        c=self.sigmoid(F.upsample_bilinear(b,scale_factor=2))*self.enc_2(c)+F.upsample_bilinear(b,scale_factor=2)
        d=self.sigmoid(F.upsample_bilinear(c,scale_factor=2))*self.enc_3(d)+F.upsample_bilinear(c,scale_factor=2)
        
        d=self.relu(self.res1(d)+d)

        d_0=self.conv_dilation_0(d)
        d_3=self.conv_dilation_3(d)
        d_6=self.conv_dilation_6(d)

        d=torch.cat([d_0,d_3,d_6],1)
        d=self.cut_dim(d)

        score=self.classifer(d)
        score=F.upsample_bilinear(score,scale_factor=4)

        return score
//  [python]
def forward(self,x):
     
        x4,x3,x2,x1,x0=self.model(x)

        a=self.enc4_1(x4)
        b=self.enc3_1(x3)
        c=self.enc2_1(x2)
        d=self.enc1_1(x1)

        b=self.sigmoid(F.interpolate(a,scale_factor=2,mode='bilinear'))*self.enc_1(b)+F.interpolate(a,scale_factor=2,mode='bilinear')
        c=self.sigmoid(F.interpolate(b,scale_factor=2,mode='bilinear'))*self.enc_2(c)+F.interpolate(b,scale_factor=2,mode='bilinear')
        d=self.sigmoid(F.interpolate(c,scale_factor=2,mode='bilinear'))*self.enc_3(d)+F.interpolate(c,scale_factor=2,mode='bilinear')
        
        d=self.relu(self.res1(d)+d)

        d_0=self.conv_dilation_0(d)
        d_3=self.conv_dilation_3(d)
        d_6=self.conv_dilation_6(d)

        d=torch.cat([d_0,d_3,d_6],1)
        d=self.cut_dim(d)

        score=self.classifer(d)
        score=F.interpolate(score,scale_factor=4,mode='bilinear')

        return score
//  [python]
def forward(self, x):
        C1, C2, C3, C4, C5 = self.resnet(x)
        P2, P3, P4, P5, P6 = self.fpn(C1, C2, C3, C4, C5)
		
        P4_x2 = F.upsample(P4, scale_factor=2)
        P5_x4 = F.upsample(P5, scale_factor=4)
        featuer_cat = torch.cat([P3, P4_x2, P5_x4],1)
        out = self.predict(featuer_cat)
        return out
//  [python]
def forward(self, x):
        C1, C2, C3, C4, C5 = self.resnet(x)
        P2, P3, P4, P5, P6 = self.fpn(C1, C2, C3, C4, C5)
		
        P4_x2 = F.interpolate(P4, scale_factor=2)
        P5_x4 = F.interpolate(P5, scale_factor=4)
        featuer_cat = torch.cat([P3, P4_x2, P5_x4],1)
        out = self.predict(featuer_cat)
        return out
//  [python]
def forward(self, x, y, w):
        x = F.upsample(x, size=16)
        x = F.upsample(x, scale_factor=2, mode='nearest')
        x = F.upsample(x, size=(20), mode='nearest')
        x = F.upsample(x, scale_factor=(4), mode='nearest')
        x = F.upsample(x, size=16, mode='linear')
        x = F.upsample(x, scale_factor=2, mode='linear')
        x = F.upsample(x, size=(24), mode='linear', align_corners=True)
        x = F.upsample(x, scale_factor=(3), mode='linear', align_corners=True)

        y = F.upsample(y, size=16)
        y = F.upsample(y, scale_factor=2, mode='nearest')
        y = F.upsample(y, size=(20,20), mode='nearest')
        y = F.upsample(y, scale_factor=(4,4), mode='nearest')
        y = F.upsample(y, size=(16,24), mode='nearest')
        y = F.upsample(y, scale_factor=(2,3), mode='nearest')
        y = F.upsample(y, size=16, mode='bilinear')
        y = F.upsample(y, scale_factor=2, mode='bilinear')
        y = F.upsample(y, size=(20,20), mode='bilinear', align_corners=False)
        y = F.upsample(y, scale_factor=(4,4), mode='bilinear', align_corners=False)
        y = F.upsample(y, size=(16,24), mode='bilinear', align_corners=True)
        y = F.upsample(y, scale_factor=(2,3), mode='bilinear', align_corners=True)
        y = F.upsample(y, size=16, mode='bicubic')
        y = F.upsample(y, scale_factor=2, mode='bicubic')
        y = F.upsample(y, size=(20,20), mode='bicubic', align_corners=False)
        y = F.upsample(y, scale_factor=(4,4), mode='bicubic', align_corners=False)
        y = F.upsample(y, size=(16,24), mode='bicubic', align_corners=True)
        y = F.upsample(y, scale_factor=(2,3), mode='bicubic', align_corners=True)

        w = F.upsample(w, scale_factor=(1.499,1.499), mode='nearest')

        return x, y, w
//  [python]
def forward(self, x, y, w):
        x = F.interpolate(x, size=16)
        x = F.interpolate(x, scale_factor=2, mode='nearest')
        x = F.interpolate(x, size=(20), mode='nearest')
        x = F.interpolate(x, scale_factor=(4), mode='nearest')
        x = F.interpolate(x, size=16, mode='linear')
        x = F.interpolate(x, scale_factor=2, mode='linear')
        x = F.interpolate(x, size=(24), mode='linear', align_corners=True)
        x = F.interpolate(x, scale_factor=(3), mode='linear', align_corners=True)

        y = F.interpolate(y, size=16)
        y = F.interpolate(y, scale_factor=2, mode='nearest')
        y = F.interpolate(y, size=(20,20), mode='nearest')
        y = F.interpolate(y, scale_factor=(4,4), mode='nearest')
        y = F.interpolate(y, size=(16,24), mode='nearest')
        y = F.interpolate(y, scale_factor=(2,3), mode='nearest')
        y = F.interpolate(y, size=16, mode='bilinear')
        y = F.interpolate(y, scale_factor=2, mode='bilinear')
        y = F.interpolate(y, size=(20,20), mode='bilinear', align_corners=False)
        y = F.interpolate(y, scale_factor=(4,4), mode='bilinear', align_corners=False)
        y = F.interpolate(y, size=(16,24), mode='bilinear', align_corners=True)
        y = F.interpolate(y, scale_factor=(2,3), mode='bilinear', align_corners=True)
        y = F.interpolate(y, size=16, mode='bicubic')
        y = F.interpolate(y, scale_factor=2, mode='bicubic')
        y = F.interpolate(y, size=(20,20), mode='bicubic', align_corners=False)
        y = F.interpolate(y, scale_factor=(4,4), mode='bicubic', align_corners=False)
        y = F.interpolate(y, size=(16,24), mode='bicubic', align_corners=True)
        y = F.interpolate(y, scale_factor=(2,3), mode='bicubic', align_corners=True)

        w = F.interpolate(w, scale_factor=(1.499,1.499), mode='nearest')

        return x, y, w
//  [python]
def forward(self, C1, C2, C3 ,C4, C5):

        p5_out = self.P5_conv1(C5)
        
        p4_out = torch.add(self.P4_conv1(C4), F.upsample_nearest(p5_out, scale_factor=2))
        p3_out = torch.add(self.P3_conv1(C3), F.upsample_nearest(p4_out, scale_factor=2))
        p2_out = torch.add(self.P2_conv1(C2), F.upsample_nearest(p3_out, scale_factor=2))

        p4_out = self.P4_conv2(p4_out)
        p3_out = self.P3_conv2(p3_out)
        p2_out = self.P2_conv2(p2_out)

        # P6 is used for the 5th anchor scale in RPN. Generated by
        # subsampling from P5 with stride of 2.
        p6_out = self.P6(p5_out)

        return p2_out, p3_out, p4_out, p5_out, p6_out
//  [python]
def forward(self, C1, C2, C3 ,C4, C5):

        p5_out = self.P5_conv1(C5)
        
        p4_out = torch.add(self.P4_conv1(C4), F.interpolate(p5_out, scale_factor=2, mode='nearest'))
        p3_out = torch.add(self.P3_conv1(C3), F.interpolate(p4_out, scale_factor=2, mode='nearest'))
        p2_out = torch.add(self.P2_conv1(C2), F.interpolate(p3_out, scale_factor=2, mode='nearest'))

        p4_out = self.P4_conv2(p4_out)
        p3_out = self.P3_conv2(p3_out)
        p2_out = self.P2_conv2(p2_out)

        # P6 is used for the 5th anchor scale in RPN. Generated by
        # subsampling from P5 with stride of 2.
        p6_out = self.P6(p5_out)

        return p2_out, p3_out, p4_out, p5_out, p6_out
//  [python]
def forward(self, x):
        out = self.conv1(self.relu(x))
        if self.droprate > 0:
            out = F.dropout(out, p=self.droprate, inplace=False, training=self.training)
        return F.upsample_nearest(out, scale_factor=2)
//  [python]
def forward(self, x):
        out = self.conv1(self.relu(x))
        if self.droprate > 0:
            out = F.dropout(out, p=self.droprate, inplace=False, training=self.training)
        return F.interpolate(out, scale_factor=2, mode='nearest')
//  [python]
def forward(self, x):
        feats = self.feats(x)
        fconn = self.fconn(feats)
        score = self.score(fconn)
        out = F.sigmoid(score)
        return F.upsample_bilinear(out, x.size()[2:])
//  [python]
def forward(self, x):
        feats = self.feats(x)
        fconn = self.fconn(feats)
        score = self.score(fconn)
        out = F.sigmoid(score)
        return F.interpolate(out, x.size()[2:], mode='bilinear')
//  [python]
def forward(self, x):
        y=self.conv3(x)
        x=self.conv1(torch.cat([x,y],1))
        return F.upsample_nearest(x, scale_factor=2)
//  [python]
def forward(self, x):
        y=self.conv3(x)
        x=self.conv1(torch.cat([x,y],1))
        return F.interpolate(x, scale_factor=2, mode='nearest')
//  [python]
def forward(self, x):
        input_size = x.size()[2:]
        feats = self.feats[0](x)
        feats3 = self.feats[1](feats)
        feats4 = self.feats[2](feats3)
        feats5 = self.feats[3](feats4)
        feats_fconn = self.fconn(feats5)

        # feats = [feats, feats3, feats4, feats5, feats_fconn]

        score_feat3 = self.score[0](feats3)
        score_feat4 = self.score[1](feats4)
        score_fconn = self.score[2](feats_fconn)

        score = F.upsample_bilinear(score_fconn, score_feat4.size()[2:])
        score += score_feat4
        score = F.upsample_bilinear(score, score_feat3.size()[2:])
        score += score_feat3

        return F.upsample_bilinear(score, input_size)
//  [python]
def forward(self, x):
        input_size = x.size()[2:]
        feats = self.feats[0](x)
        feats3 = self.feats[1](feats)
        feats4 = self.feats[2](feats3)
        feats5 = self.feats[3](feats4)
        feats_fconn = self.fconn(feats5)

        # feats = [feats, feats3, feats4, feats5, feats_fconn]

        score_feat3 = self.score[0](feats3)
        score_feat4 = self.score[1](feats4)
        score_fconn = self.score[2](feats_fconn)

        score = F.interpolate(score_fconn, size=score_feat4.size()[2:], mode='bilinear', align_corners=True)
        score += score_feat4
        score = F.interpolate(score, size=score_feat3.size()[2:], mode='bilinear', align_corners=True)
        score += score_feat3

        return F.interpolate(score, size=input_size, mode='bilinear', align_corners=True)
//  [python]
def forward(self, lr_ms, hr_pan, hr_ms, fuse_ms):
        down_hr_pan = torch.nn.functional.upsample(hr_pan, scale_factor=1/self.scale, mode='area')

        M_ = torch.zeros(hr_pan.shape).cuda()
        G_ = torch.zeros(hr_pan.shape).cuda()
        for j in range(lr_ms.shape[0]):
            down_hr_pan_array = down_hr_pan[j].view(-1)
            down_hr_pan_array = fromDlpack(to_dlpack(down_hr_pan_array))
            for i in range(lr_ms.shape[1]):
                lr_ms_band = lr_ms[j,i,:,:].squeeze(0)
                lr_ms_band_array = lr_ms_band.view(-1)
                lr_ms_band_array = fromDlpack(to_dlpack(lr_ms_band_array))
                if i == 0:
                    A = cp.vstack([lr_ms_band_array**0, lr_ms_band_array**1])
                else:
                    A = cp.vstack([A,  lr_ms_band_array**1])
            
            sol, r, rank, s = cp.linalg.lstsq(A.T,down_hr_pan_array)
          
            sol = from_dlpack(toDlpack(sol))
            
            for i in range(hr_ms.shape[1]):
                M_[j] += hr_ms[j,i,:,:] * sol[i+1]
                G_[j] += fuse_ms[j,i,:,:] * sol[i+1]
            M_[j] += sol[0]
            G_[j] += sol[0] 
        
        mean_filter = kornia.filters.BoxBlur((31,31))
        e = torch.Tensor([1e-10]).cuda()
        r = 4
        a = 1 
        mean_M_ = mean_filter(M_)
        mean_P =  mean_filter(hr_pan) 
        mean_M_xP = mean_filter(M_*hr_pan)
        cov_M_xP = mean_M_xP - mean_M_*mean_P
        mean_M_xM_ = mean_filter(M_*M_)
        std_M_ = torch.sqrt(torch.abs(mean_M_xM_ - mean_M_*mean_M_) + e)
        mean_PxP = mean_filter(hr_pan*hr_pan)
        std_P = torch.sqrt(torch.abs(mean_PxP - mean_P*mean_P) + e)
        corr_M_xP = cov_M_xP / (std_M_*std_P)
        S = corr_M_xP**r
        loss_c = self.L1loss(fuse_ms*S, hr_ms*S)

        grad_P = (hr_pan - mean_P) / std_P
        mean_G_ = mean_filter(G_)
        mean_G_xG_ = mean_filter(G_*G_)
        std_G_ = torch.sqrt(torch.abs(mean_G_xG_ - mean_G_*mean_G_) + e)
        grad_G_ = (G_ - mean_G_) / std_G_
        loss_a = self.L1loss(grad_G_*(2-S), grad_P*(2-S))

        loss = loss_c + a*loss_a
        return loss
//  [python]
def forward(self, lr_ms, hr_pan, hr_ms, fuse_ms):
        down_hr_pan = torch.nn.functional.interpolate(hr_pan, scale_factor=1/self.scale, mode='area')

        M_ = torch.zeros(hr_pan.shape).cuda()
        G_ = torch.zeros(hr_pan.shape).cuda()
        for j in range(lr_ms.shape[0]):
            down_hr_pan_array = down_hr_pan[j].view(-1)
            down_hr_pan_array = fromDlpack(to_dlpack(down_hr_pan_array))
            for i in range(lr_ms.shape[1]):
                lr_ms_band = lr_ms[j,i,:,:].squeeze(0)
                lr_ms_band_array = lr_ms_band.view(-1)
                lr_ms_band_array = fromDlpack(to_dlpack(lr_ms_band_array))
                if i == 0:
                    A = cp.vstack([lr_ms_band_array**0, lr_ms_band_array**1])
                else:
                    A = cp.vstack([A,  lr_ms_band_array**1])
            
            sol, r, rank, s = cp.linalg.lstsq(A.T,down_hr_pan_array)
          
            sol = from_dlpack(toDlpack(sol))
            
            for i in range(hr_ms.shape[1]):
                M_[j] += hr_ms[j,i,:,:] * sol[i+1]
                G_[j] += fuse_ms[j,i,:,:] * sol[i+1]
            M_[j] += sol[0]
            G_[j] += sol[0] 
        
        mean_filter = kornia.filters.BoxBlur((31,31))
        e = torch.Tensor([1e-10]).cuda()
        r = 4
        a = 1 
        mean_M_ = mean_filter(M_)
        mean_P =  mean_filter(hr_pan) 
        mean_M_xP = mean_filter(M_*hr_pan)
        cov_M_xP = mean_M_xP - mean_M_*mean_P
        mean_M_xM_ = mean_filter(M_*M_)
        std_M_ = torch.sqrt(torch.abs(mean_M_xM_ - mean_M_*mean_M_) + e)
        mean_PxP = mean_filter(hr_pan*hr_pan)
        std_P = torch.sqrt(torch.abs(mean_PxP - mean_P*mean_P) + e)
        corr_M_xP = cov_M_xP / (std_M_*std_P)
        S = corr_M_xP**r
        loss_c = self.L1loss(fuse_ms*S, hr_ms*S)

        grad_P = (hr_pan - mean_P) / std_P
        mean_G_ = mean_filter(G_)
        mean_G_xG_ = mean_filter(G_*G_)
        std_G_ = torch.sqrt(torch.abs(mean_G_xG_ - mean_G_*mean_G_) + e)
        grad_G_ = (G_ - mean_G_) / std_G_
        loss_a = self.L1loss(grad_G_*(2-S), grad_P*(2-S))

        loss = loss_c + a*loss_a
        return loss
//  [python]
def forward(self, features1, features2):
        # unify the size of width and height
        B, C, H, W = features1.size()
        _, M, AH, AW = features2.size()

        # match size
        if AH != H or AW != W:
            features2 = F.upsample_bilinear(features2, size=(H, W))

        # essential_matrix: (B, M, C) -> (B, M * C)
        essential_matrix = (torch.einsum('imjk,injk->imn', (features2, features1)) / float(H * W)).view(B, -1)
        # nornalize
        essential_matrix = torch.sign(essential_matrix) * torch.sqrt(torch.abs(essential_matrix) + self.epsilon)
        essential_matrix = F.normalize(essential_matrix, dim=-1)

        return essential_matrix
//  [python]
def forward(self, features1, features2):
        # unify the size of width and height
        B, C, H, W = features1.size()
        _, M, AH, AW = features2.size()

        # match size
        if AH != H or AW != W:
            features2 = F.interpolate(features2, size=(H, W), mode='bilinear')

        # essential_matrix: (B, M, C) -> (B, M * C)
        essential_matrix = (torch.einsum('imjk,injk->imn', (features2, features1)) / float(H * W)).view(B, -1)
        # nornalize
        essential_matrix = torch.sign(essential_matrix) * torch.sqrt(torch.abs(essential_matrix) + self.epsilon)
        essential_matrix = F.normalize(essential_matrix, dim=-1)

        return essential_matrix
//  [python]
def forward(self,x,alpha=0.5,downsamples=0,embedding=None,useStylePredictor=True,detach=True):
        # augments a batch of images with style randomization
        # x: B x C x H x W image tensor
        # alpha: float in [0,1], controls interpolation between random style and original style
        # downsamples: int, number of times to downsample by factor of 2 before applying style transfer
        # embedding: B x 100 tensor, or None. Use this embedding if provided.
        # useStylePredictor: bool. If True, we use the inception based style predictor to compute the original style embedding for the input image, and use that for interpolation. If False, we use the mean ImageNet embedding instead, which is slightly faster.
        # detach: bool. If true, detach the augmented image tensor to prevent backpropagation into self.ghiasi and self.stylePredictor

        # style embedding for when alpha=0:
        base = self.stylePredictor(x) if useStylePredictor else self.imagenet_embedding

        if downsamples:
            assert(x.size(2) % 2**downsamples == 0)
            assert(x.size(3) % 2**downsamples == 0)
            for i in range(downsamples):
                x = nn.functional.avg_pool2d(x,2)

        if embedding is None:
            # sample a random embedding
            embedding = self.sample_embedding(x.size(0))
        # interpolate style embeddings:
        embedding = alpha*embedding + (1-alpha)*base
        
        restyled = self.ghiasi(x,embedding)

        if downsamples:
            restyled = nn.functional.upsample(restyled,scale_factor=2**downsamples,mode='bilinear')
        
        if detach:
            return restyled.detach()
        else:
            return restyled
//  [python]
def forward(self,x,alpha=0.5,downsamples=0,embedding=None,useStylePredictor=True,detach=True):
        # augments a batch of images with style randomization
        # x: B x C x H x W image tensor
        # alpha: float in [0,1], controls interpolation between random style and original style
        # downsamples: int, number of times to downsample by factor of 2 before applying style transfer
        # embedding: B x 100 tensor, or None. Use this embedding if provided.
        # useStylePredictor: bool. If True, we use the inception based style predictor to compute the original style embedding for the input image, and use that for interpolation. If False, we use the mean ImageNet embedding instead, which is slightly faster.
        # detach: bool. If true, detach the augmented image tensor to prevent backpropagation into self.ghiasi and self.stylePredictor

        # style embedding for when alpha=0:
        base = self.stylePredictor(x) if useStylePredictor else self.imagenet_embedding

        if downsamples:
            assert(x.size(2) % 2**downsamples == 0)
            assert(x.size(3) % 2**downsamples == 0)
            for i in range(downsamples):
                x = nn.functional.avg_pool2d(x,2)

        if embedding is None:
            # sample a random embedding
            embedding = self.sample_embedding(x.size(0))
        # interpolate style embeddings:
        embedding = alpha*embedding + (1-alpha)*base
        
        restyled = self.ghiasi(x,embedding)

        if downsamples:
            restyled = nn.functional.interpolate(restyled,scale_factor=2**downsamples,mode='bilinear')
        
        if detach:
            return restyled.detach()
        else:
            return restyled