jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def rigid_transform_Kabsch_3D_torch(
        A: Tensor, # shape (3, N)
        B: Tensor, # ref shape (3, N)
):
    # R = 3x3 rotation matrix, t = 3x1 column vector
    assert A.shape[1] == B.shape[1]
    num_rows, num_cols = A.shape
    if num_rows != 3:
        raise Exception(f"matrix A is not 3xN, it is {num_rows}x{num_cols}")
    num_rows, num_cols = B.shape
    if num_rows != 3:
        raise Exception(f"matrix B is not 3xN, it is {num_rows}x{num_cols}")

    # find mean column wise: 3 x 1
    centroid_A = torch.mean(A, dim = 1, keepdim = True)
    centroid_B = torch.mean(B, dim = 1, keepdim = True)

    # subtract mean
    Am = A - centroid_A
    Bm = B - centroid_B

    H = Am @ Bm.T

    # find rotation
    U, S, Vt = torch.linalg.svd(H)

    # Compared with SVD_rigid_transform_torch‘s implementation R^T because (N,3_a,3_b) \matmul (N, 3_b) -> (N, 3_b) @ (N, 3_b, 3_a)
    # Here the coordinates (3, N) so we can use R, so (N, 3_a, 3_b) @ (N, 3_b)
    R = Vt.T @ U.T # R \equiv (N, 3_a, 3_b)
    if torch.linalg.det(R) < 0:
        SS = torch.diag(torch.tensor([1.,1.,-1.], device = A.device))
        R = (Vt.T @ SS) @ U.T
    assert math.fabs(torch.linalg.det(R) - 1) < 3e-3

    t = -R @ centroid_A + centroid_B
    return R, t
//  [python]
def align(
            self,
            ref_coords: Optional[Tensor] = None,
            coords: Optional[Tensor] = None,
    ):
        self._clear()
        if ref_coords is None or coords is None:
            raise AttributeError(f":cls:`{self.__class__.__name__}` :attr:`ref_coords` or `coords` "
                                 f" is None. Please set them by calling the :func:`align`.")
        ref_shape = ref_coords.shape
        shape = coords.shape
        if ref_shape != shape:
            raise ValueError("Input `ref_coords` and `coords` must be the same shape, "
                             f"but got `ref_coords` shape {ref_shape}, `coords` shape {shape}")
        if ref_shape[1] == 3 and shape[1] == 3:
            raise ValueError("The two dimension of input `ref_coords` and `coords` must be 3 for "
                             f"Euclidean space, but get `ref_coords` shape {ref_shape}, `coords` shape {shape}")
        self._ref_coords = ref_coords # shape (N, 3)
        self._coords = coords # shape (N, 3)
        self._init_rmsd = self.calc_rmsd(ref_coords, coords)
        # get centroids for ref and coords and center on the centroids
        # centroid shape (1, 3)
        centroid_ref = torch.mean(ref_coords, dim = 0, keepdim = True)
        centroid_coords = torch.mean(coords, dim = 0, keepdim = True)
        ref_coords = ref_coords - centroid_ref
        coords = coords - centroid_coords
        # get correlation matrix
        C = coords.T @ ref_coords
        # SVG decomposition
        U, S, V_T = torch.svd(C)
        # find rotation matrix
        self._rot = (V_T.T @ U.T).T
        # reflection
        if torch.det(self._rot) < 0:
            V_T[2] = -V_T[2]
            self._rot = (V_T.T @ U.T).T
        # translation matrix shape (1, 3)
        self._trans = centroid_ref - centroid_coords @ self.rot
        self._trans_coords = self._coords @ self.rot + self.trans
        self._rmsd = self.calc_rmsd(self._ref_coords, self._trans_coords)
//  [python]
def align(
            self,
            ref_coords: Optional[Tensor] = None,
            coords: Optional[Tensor] = None,
    ):
        self._clear()
        if ref_coords is None or coords is None:
            raise AttributeError(f":cls:`{self.__class__.__name__}` :attr:`ref_coords` or `coords` "
                                 f" is None. Please set them by calling the :func:`align`.")
        ref_shape = ref_coords.shape
        shape = coords.shape
        if ref_shape != shape:
            raise ValueError("Input `ref_coords` and `coords` must be the same shape, "
                             f"but got `ref_coords` shape {ref_shape}, `coords` shape {shape}")
        if ref_shape[1] == 3 and shape[1] == 3:
            raise ValueError("The two dimension of input `ref_coords` and `coords` must be 3 for "
                             f"Euclidean space, but get `ref_coords` shape {ref_shape}, `coords` shape {shape}")
        self._ref_coords = ref_coords # shape (N, 3)
        self._coords = coords # shape (N, 3)
        self._init_rmsd = self.calc_rmsd(ref_coords, coords)
        # get centroids for ref and coords and center on the centroids
        # centroid shape (1, 3)
        centroid_ref = torch.mean(ref_coords, dim = 0, keepdim = True)
        centroid_coords = torch.mean(coords, dim = 0, keepdim = True)
        ref_coords = ref_coords - centroid_ref
        coords = coords - centroid_coords
        # get correlation matrix
        C = coords.T @ ref_coords
        # SVG decomposition
        U, S, V_T = torch.linalg.svd(C)
        # find rotation matrix
        self._rot = (V_T.T @ U.T).T
        # reflection
        if torch.linalg.det(self._rot) < 0:
            V_T[2] = -V_T[2]
            self._rot = (V_T.T @ U.T).T
        # translation matrix shape (1, 3)
        self._trans = centroid_ref - centroid_coords @ self.rot
        self._trans_coords = self._coords @ self.rot + self.trans
        self._rmsd = self.calc_rmsd(self._ref_coords, self._trans_coords)
//  [python]
def solve(task=None, L=None, batch_shape=(), variable=True, **kwargs):
    N = int(np.max(np.roots([1, 2, -L])))
    assert N**2 + N + N <= L
   
    if 'fixed' in task: variable = False
    # sample orthogonal linear system
    if variable:
        A = torch.randn(batch_shape + (N,N))    # (B N N)
    else:
        A = torch.randn((N,N), generator=torch.Generator().manual_seed(0))     # (N N)
        A = A.view((1,)*len(batch_shape) + (N,N)).expand(batch_shape + (N,N))  # (B N N)
    
    U, S, Vh = torch.svd(A)
    A = U.matmul(torch.diag_embed(torch.ones(U.shape[1]))).matmul(Vh)                                           # A.A^T == I
    
    X = F.normalize(torch.randn(batch_shape + (N,1)), dim=-2)  # (B N 1)
    B = A.matmul(X)  # AX == B                                 # (B N 1)
    
    AB = torch.cat((A, B), dim=-1)                             # (B N N+1)
    x = AB.view(*AB.shape[:-2], -1)                            # (B N^2+N)
    x = F.pad(x, (0,L-x.size(-1))).unsqueeze(-1)               # (B L 1)
    y = X
    return concat_pos(x), y
//  [python]
def solve(task=None, L=None, batch_shape=(), variable=True, **kwargs):
    N = int(np.max(np.roots([1, 2, -L])))
    assert N**2 + N + N <= L
   
    if 'fixed' in task: variable = False
    # sample orthogonal linear system
    if variable:
        A = torch.randn(batch_shape + (N,N))    # (B N N)
    else:
        A = torch.randn((N,N), generator=torch.Generator().manual_seed(0))     # (N N)
        A = A.view((1,)*len(batch_shape) + (N,N)).expand(batch_shape + (N,N))  # (B N N)
    
    U, _, Vh = torch.linalg.svd(A, full_matrices=False)
    A = U.matmul(Vh)                                           # A.A^T == I
    
    X = F.normalize(torch.randn(batch_shape + (N,1)), dim=-2)  # (B N 1)
    B = A.matmul(X)  # AX == B                                 # (B N 1)
    
    AB = torch.cat((A, B), dim=-1)                             # (B N N+1)
    x = AB.view(*AB.shape[:-2], -1)                            # (B N^2+N)
    x = F.pad(x, (0,L-x.size(-1))).unsqueeze(-1)               # (B L 1)
    y = X
    return concat_pos(x), y
//  [python]
def forward(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.pool 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.pool(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) + EPSILON)

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

        if self.training:
            fake_att = torch.zeros_like(attentions).uniform_(0, 2)
        else:
            fake_att = torch.ones_like(attentions)
        counterfactual_feature = (torch.einsum('imjk,injk->imn', (fake_att, features)) / float(H * W)).view(B, -1)

        counterfactual_feature = torch.sign(counterfactual_feature) * torch.sqrt(torch.abs(counterfactual_feature) + EPSILON)

        counterfactual_feature = F.normalize(counterfactual_feature, dim=-1)
        return feature_matrix, counterfactual_feature
//  [python]
def forward(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.pool 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.pool(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) + EPSILON)

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

        if self.training:
            fake_att = torch.zeros_like(attentions).uniform_(0, 2)
        else:
            fake_att = torch.ones_like(attentions)
        counterfactual_feature = (torch.einsum('imjk,injk->imn', (fake_att, features)) / float(H * W)).view(B, -1)

        counterfactual_feature = torch.sign(counterfactual_feature) * torch.sqrt(torch.abs(counterfactual_feature) + EPSILON)

        counterfactual_feature = F.normalize(counterfactual_feature, dim=-1)
        return feature_matrix, counterfactual_feature
//  [python]
def forward(self, x):
        out = self.conv1(self.relu(self.bn1(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(self.bn1(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 train_new(args, model_s, train_loader, device):
    def calculate_linear_regression(X, Y, singular_value_threshold):
        # Calculate the means of X and Y
        X_mean = torch.mean(X, axis=0)
        Y_mean = torch.mean(Y, axis=0)
        # Center X and Y
        X_centered = X - X_mean
        Y_centered = Y - Y_mean        
        # Compute the SVD of centered X
        U, S, Vt = torch.svd(X_centered)
        # Apply the threshold to the singular values
        S_inv = torch.where(S > singular_value_threshold, 1.0/S, 0.0)
        num_fulfilled = torch.sum(S_inv.nonzero()).item()
        # Compute the pseudoinverse of centered X using the SVD
        X_centered_pinv = torch.matmul(Vt.T, torch.matmul(torch.diag(S_inv), U.T))
        # Compute A using least squares
        A = torch.matmul(X_centered_pinv, Y_centered)
        # Compute bias term b
        b = Y_mean - torch.matmul(X_mean, A)
        return A, b, num_fulfilled


    # Assuming X and Y are your data
    # my batch size is the full dataset
    for _, (data, target, wanted_token) in enumerate(train_loader):
        data, target, wanted_token = data.to(device), target.to(device), wanted_token.to(device)
        singular_value_threshold = 0.0
        A, b, dim_s_inv = calculate_linear_regression(data, target, singular_value_threshold)

        with torch.no_grad():
            model_s.fc.weight[...] = A.T
            model_s.fc.bias[...] = b

        display_name = args["display_name"]
        args["test_results"] = os.path.join(args["experiment_data_dir"],
                               f"linear_model_train_{display_name}_full_math_results.csv")
        
        #train_loss, train_acc, train_gt_acc = test(args, model_s, device, train_loader)
        saved_model_path = os.path.join(args["experiment_data_dir"],f'model_trained_{display_name}_math_sol.pth')
        torch.save(model_s.state_dict(), saved_model_path)
        return
//  [python]
def train_new(args, model_s, train_loader, device):
    def calculate_linear_regression(X, Y, singular_value_threshold):
        # Calculate the means of X and Y
        X_mean = torch.mean(X, axis=0)
        Y_mean = torch.mean(Y, axis=0)
        # Center X and Y
        X_centered = X - X_mean
        Y_centered = Y - Y_mean        
        # Compute the SVD of centered X
        U, S, Vt = torch.linalg.svd(X_centered, full_matrices=False)
        # Apply the threshold to the singular values
        S_inv = torch.where(S > singular_value_threshold, 1.0/S, 0.0)
        num_fulfilled = torch.sum(S_inv.nonzero()).item()
        # Compute the pseudoinverse of centered X using the SVD
        X_centered_pinv = torch.matmul(Vt.T, torch.matmul(torch.diag(S_inv), U.T))
        # Compute A using least squares
        A = torch.matmul(X_centered_pinv, Y_centered)
        # Compute bias term b
        b = Y_mean - torch.matmul(X_mean, A)
        return A, b, num_fulfilled


    # Assuming X and Y are your data
    # my batch size is the full dataset
    for _, (data, target, wanted_token) in enumerate(train_loader):
        data, target, wanted_token = data.to(device), target.to(device), wanted_token.to(device)
        singular_value_threshold = 0.0
        A, b, dim_s_inv = calculate_linear_regression(data, target, singular_value_threshold)

        with torch.no_grad():
            model_s.fc.weight[...] = A.T
            model_s.fc.bias[...] = b

        display_name = args["display_name"]
        args["test_results"] = os.path.join(args["experiment_data_dir"],
                               f"linear_model_train_{display_name}_full_math_results.csv")
        
        #train_loss, train_acc, train_gt_acc = test(args, model_s, device, train_loader)
        saved_model_path = os.path.join(args["experiment_data_dir"],f'model_trained_{display_name}_math_sol.pth')
        torch.save(model_s.state_dict(), saved_model_path)
        return
//  [python]
def batch_augment(images, attention_map, mode='crop', theta=0.5, padding_ratio=0.1):
    batches, _, imgH, imgW = images.size()

    if mode == 'crop':
        crop_images = []
        for batch_index in range(batches):
            atten_map = attention_map[batch_index:batch_index + 1]
            if isinstance(theta, tuple):
                theta_c = random.uniform(*theta) * atten_map.max()
            else:
                theta_c = theta * atten_map.max()

            crop_mask = F.upsample_bilinear(atten_map, size=(imgH, imgW)) >= theta_c
            nonzero_indices = torch.nonzero(crop_mask[0, 0, ...])
            height_min = max(int(nonzero_indices[:, 0].min().item() - padding_ratio * imgH), 0)
            height_max = min(int(nonzero_indices[:, 0].max().item() + padding_ratio * imgH), imgH)
            width_min = max(int(nonzero_indices[:, 1].min().item() - padding_ratio * imgW), 0)
            width_max = min(int(nonzero_indices[:, 1].max().item() + padding_ratio * imgW), imgW)

            crop_images.append(
                F.upsample_bilinear(images[batch_index:batch_index + 1, :, height_min:height_max, width_min:width_max],
                                    size=(imgH, imgW)))
        crop_images = torch.cat(crop_images, dim=0)
        return crop_images

    elif mode == 'drop':
        drop_masks = []
        for batch_index in range(batches):
            atten_map = attention_map[batch_index:batch_index + 1]
            if isinstance(theta, tuple):
                theta_d = random.uniform(*theta) * atten_map.max()
            else:
                theta_d = theta * atten_map.max()

            drop_masks.append(F.upsample_bilinear(atten_map, size=(imgH, imgW)) < theta_d)
        drop_masks = torch.cat(drop_masks, dim=0)
        drop_images = images * drop_masks.float()
        return drop_images

    else:
        raise ValueError('Expected mode in [\'crop\', \'drop\'], but received unsupported augmentation method %s' % mode)
//  [python]
def batch_augment(images, attention_map, mode='crop', theta=0.5, padding_ratio=0.1):
    batches, _, imgH, imgW = images.size()

    if mode == 'crop':
        crop_images = []
        for batch_index in range(batches):
            atten_map = attention_map[batch_index:batch_index + 1]
            if isinstance(theta, tuple):
                theta_c = random.uniform(*theta) * atten_map.max()
            else:
                theta_c = theta * atten_map.max()

            crop_mask = F.interpolate(atten_map, size=(imgH, imgW), mode='bilinear') >= theta_c
            nonzero_indices = torch.nonzero(crop_mask[0, 0, ...])
            height_min = max(int(nonzero_indices[:, 0].min().item() - padding_ratio * imgH), 0)
            height_max = min(int(nonzero_indices[:, 0].max().item() + padding_ratio * imgH), imgH)
            width_min = max(int(nonzero_indices[:, 1].min().item() - padding_ratio * imgW), 0)
            width_max = min(int(nonzero_indices[:, 1].max().item() + padding_ratio * imgW), imgW)

            crop_images.append(
                F.interpolate(images[batch_index:batch_index + 1, :, height_min:height_max, width_min:width_max],
                                    size=(imgH, imgW), mode='bilinear'))
        crop_images = torch.cat(crop_images, dim=0)
        return crop_images

    elif mode == 'drop':
        drop_masks = []
        for batch_index in range(batches):
            atten_map = attention_map[batch_index:batch_index + 1]
            if isinstance(theta, tuple):
                theta_d = random.uniform(*theta) * atten_map.max()
            else:
                theta_d = theta * atten_map.max()

            drop_masks.append(F.interpolate(atten_map, size=(imgH, imgW), mode='bilinear') < theta_d)
        drop_masks = torch.cat(drop_masks, dim=0)
        drop_images = images * drop_masks.float()
        return drop_images

    else:
        raise ValueError('Expected mode in [\'crop\', \'drop\'], but received unsupported augmentation method %s' % mode)
//  [python]
def forward(self, x, upsample=None):
        x = self._pre_process(x)
        x = self.model(x)
        if upsample is not None:
            x = torch.nn.functional.upsample_bilinear(x, size=upsample)
        return x
//  [python]
def forward(self, x, upsample=None):
        x = self._pre_process(x)
        x = self.model(x)
        if upsample is not None:
            x = torch.nn.functional.interpolate(x, size=upsample, mode='bilinear')
        return x
//  [python]
def forward(self, X):
        """Forward pass of the network.
        """
        N = X.size()[0]
        X1 = self.features1(X)
        H = X1.size()[2]
        W = X1.size()[3]
        assert X1.size()[1] == 512
        X2 = self.features2(X)
        H2 = X2.size()[2]
        W2 = X2.size()[3]
        assert X2.size()[1] == 128

        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 = torch.nn.functional.normalize(X)
        X = self.fc(X)
        assert X.size() == (N, 2)

        if self.config.std_modeling:
            mean = X [:, 0]
            t = X [:, 1]
            var = nn.functional.softplus(t)
            return mean, var
        else:
            return X
//  [python]
def forward(self, X):
        """Forward pass of the network.
        """
        N = X.size()[0]
        X1 = self.features1(X)
        H = X1.size()[2]
        W = X1.size()[3]
        assert X1.size()[1] == 512
        X2 = self.features2(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 = torch.nn.functional.normalize(X)
        X = self.fc(X)
        assert X.size() == (N, 2)

        if self.config.std_modeling:
            mean = X [:, 0]
            t = X [:, 1]
            var = nn.functional.softplus(t)
            return mean, var
        else:
            return X
//  [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.refine1(self.gcn1(fm4))
        gcfm2 = self.refine2(self.gcn2(fm3))
        gcfm3 = self.refine3(self.gcn3(fm2))
        gcfm4 = self.refine4(self.gcn4(pool_x))
        gcfm5 = self.refine5(self.gcn5(conv_x))

        fs1 = self.refine6(F.upsample_bilinear(gcfm1, fm3.size()[2:]) + gcfm2)
        fs2 = self.refine7(F.upsample_bilinear(fs1, fm2.size()[2:]) + gcfm3)
        fs3 = self.refine8(F.upsample_bilinear(fs2, pool_x.size()[2:]) + gcfm4)
        fs4 = self.refine9(F.upsample_bilinear(fs3, conv_x.size()[2:]) + gcfm5)
        out = self.refine10(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.refine1(self.gcn1(fm4))
        gcfm2 = self.refine2(self.gcn2(fm3))
        gcfm3 = self.refine3(self.gcn3(fm2))
        gcfm4 = self.refine4(self.gcn4(pool_x))
        gcfm5 = self.refine5(self.gcn5(conv_x))

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

        return out
//  [python]
def forward(self, error, top_down_state, state):
        if error is None:  # we just started
            error = Variable(torch.zeros(self.error_init_size))
        model_input = error
        if top_down_state is not None:
            model_input = torch.cat((error, f.upsample_nearest(top_down_state, scale_factor=2)), 1)
        return self.memory(model_input, state)
//  [python]
def forward(self, error, top_down_state, state):
        if error is None:  # we just started
            error = Variable(torch.zeros(self.error_init_size))
        model_input = error
        if top_down_state is not None:
            model_input = torch.cat((error, f.interpolate(top_down_state, scale_factor=2, mode='nearest')), 1)
        return self.memory(model_input, state)
//  [python]
def forward(self, feats):
        h, w = feats.size(2), feats.size(3)
        priors = [
            F.upsample(input=stage(feats), size=(h, w), mode='bilinear')
            for stage in self.stages
        ] + [feats]
        bottle = self.bottleneck(torch.cat(priors, 1))
        return self.relu(bottle)
//  [python]
def forward(self, feats):
        h, w = feats.size(2), feats.size(3)
        priors = [
            F.interpolate(input=stage(feats), size=(h, w), mode='bilinear')
            for stage in self.stages
        ] + [feats]
        bottle = self.bottleneck(torch.cat(priors, 1))
        return self.relu(bottle)
//  [python]
def forward(self, inputs):
        visual_feats = inputs["visual"]
        language_dict_features = inputs["lang"]

        next_x = []
        for level, feature in enumerate(visual_feats):

            conv_args = dict()
            if self.offset is not None:
                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](visual_feats[level - 1], **conv_args))
            if level < len(visual_feats) - 1:
                temp_fea.append(
                    F.upsample_bilinear(
                        self.DyConv[0](visual_feats[level + 1], **conv_args), size=[feature.size(2), feature.size(3)]
                    )
                )
            mean_fea = torch.mean(torch.stack(temp_fea), dim=0, keepdim=False)

            if self.AttnConv is not None:
                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.append(mean_fea)

        next_x = [self.relu(item) for item in next_x]

        features_dict = {"visual": next_x, "lang": language_dict_features}

        return features_dict
//  [python]
def forward(self, inputs):
        visual_feats = inputs["visual"]
        language_dict_features = inputs["lang"]

        next_x = []
        for level, feature in enumerate(visual_feats):

            conv_args = dict()
            if self.offset is not None:
                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](visual_feats[level - 1], **conv_args))
            if level < len(visual_feats) - 1:
                temp_fea.append(
                    F.interpolate(
                        self.DyConv[0](visual_feats[level + 1], **conv_args), size=[feature.size(2), feature.size(3)], mode='bilinear'
                    )
                )
            mean_fea = torch.mean(torch.stack(temp_fea), dim=0, keepdim=False)

            if self.AttnConv is not None:
                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.append(mean_fea)

        next_x = [self.relu(item) for item in next_x]

        features_dict = {"visual": next_x, "lang": language_dict_features}

        return features_dict
//  [python]
def forward(self, x):
        # save for combining later with output
        residual = x

        if self.combine == 'concat':
            residual = F.avg_pool2d(residual, kernel_size=3, stride=2, padding=1)
            if self.up is True:
                residual = F.upsample(residual, scale_factor=2, mode="bilinear")

        out = self.g_conv_1x1_compress(x)
        out = channel_shuffle(out, self.groups)

        out = self.depthwise_conv3x3(out)
        out = self.bn_after_depthwise(out)

        out = self.g_conv_1x1_expand(out)

        out = self._combine_func(residual, out)
        return F.relu(out)
//  [python]
def forward(self, x):
        # save for combining later with output
        residual = x

        if self.combine == 'concat':
            residual = F.avg_pool2d(residual, kernel_size=3, stride=2, padding=1)
            if self.up is True:
                residual = F.interpolate(residual, scale_factor=2, mode="bilinear")

        out = self.g_conv_1x1_compress(x)
        out = channel_shuffle(out, self.groups)

        out = self.depthwise_conv3x3(out)
        out = self.bn_after_depthwise(out)

        out = self.g_conv_1x1_expand(out)

        out = self._combine_func(residual, out)
        return F.relu(out)
//  [python]
def run_on_tensor(weights, ip_tensor, num_classes):
    i3d = InceptionI3d(400, in_channels=3)
    # i3d.load_state_dict(torch.load('models/rgb_imagenet.pt'))

    i3d.replace_logits(num_classes)
    i3d.load_state_dict(torch.load(weights))  # nslt_2000_000700.pt nslt_1000_010800 nslt_300_005100.pt(best_results)  nslt_300_005500.pt(results_reported) nslt_2000_011400
    i3d.cuda()
    i3d = nn.DataParallel(i3d)
    i3d.eval()

    t = ip_tensor.shape[2]
    ip_tensor.cuda()
    per_frame_logits = i3d(ip_tensor)

    predictions = F.upsample(per_frame_logits, t, mode='linear')

    predictions = predictions.transpose(2, 1)
    out_labels = np.argsort(predictions.cpu().detach().numpy()[0])

    arr = predictions.cpu().detach().numpy()[0,:,0].T

    plt.plot(range(len(arr)), F.softmax(torch.from_numpy(arr), dim=0).numpy())
    plt.show()

    return out_labels
//  [python]
def run_on_tensor(weights, ip_tensor, num_classes):
    i3d = InceptionI3d(400, in_channels=3)
    # i3d.load_state_dict(torch.load('models/rgb_imagenet.pt'))

    i3d.replace_logits(num_classes)
    i3d.load_state_dict(torch.load(weights))  # nslt_2000_000700.pt nslt_1000_010800 nslt_300_005100.pt(best_results)  nslt_300_005500.pt(results_reported) nslt_2000_011400
    i3d.cuda()
    i3d = nn.DataParallel(i3d)
    i3d.eval()

    t = ip_tensor.shape[2]
    ip_tensor.cuda()
    per_frame_logits = i3d(ip_tensor)

    predictions = F.interpolate(per_frame_logits, t, mode='linear')

    predictions = predictions.transpose(2, 1)
    out_labels = np.argsort(predictions.cpu().detach().numpy()[0])

    arr = predictions.cpu().detach().numpy()[0,:,0].T

    plt.plot(range(len(arr)), F.softmax(torch.from_numpy(arr), dim=0).numpy())
    plt.show()

    return out_labels
//  [python]
def create_gradient_masks(input_img, cls_model):

    # input preprocessing
    input_img.requires_grad = True
    batch_size = input_img.size(0)
    width = input_img.size(2)
    height = input_img.size(3)
    cls_model.eval()
    batch_masks = []
    pred_cls, k = cls_model(input_img)

    for b in range(batch_size):
        target_labels = 1  # snow part
        grads = torch.autograd.grad(pred_cls[b][target_labels], k, retain_graph=True)[0]
        one_grad = grads[b]
        a_k = torch.mean(one_grad, dim=(1, 2), keepdim=True)
        mask = torch.sum(a_k * k[b], dim=0).to(device)
        out = mask

        # normalization
        out = (out + torch.abs(out)) / 2
        if out.sum() == 0:  # no division by zero
            out = out
        else:
            out = out / torch.max(out)  # 0 ~ 1 scaling

        # up-sampling
        out = F.upsample_bilinear(out.unsqueeze(0).unsqueeze(0), [width, height])  # 4D로 바꿈

        batch_masks.append(out)
    batch_masks = torch.stack(batch_masks, dim=0).squeeze(1)  # [B, 1, 1, 240, 240] --> [B, 1, 240, 240]
    return batch_masks
//  [python]
def create_gradient_masks(input_img, cls_model):

    # input preprocessing
    input_img.requires_grad = True
    batch_size = input_img.size(0)
    width = input_img.size(2)
    height = input_img.size(3)
    cls_model.eval()
    batch_masks = []
    pred_cls, k = cls_model(input_img)

    for b in range(batch_size):
        target_labels = 1  # snow part
        grads = torch.autograd.grad(pred_cls[b][target_labels], k, retain_graph=True)[0]
        one_grad = grads[b]
        a_k = torch.mean(one_grad, dim=(1, 2), keepdim=True)
        mask = torch.sum(a_k * k[b], dim=0).to(device)
        out = mask

        # normalization
        out = (out + torch.abs(out)) / 2
        if out.sum() == 0:  # no division by zero
            out = out
        else:
            out = out / torch.max(out)  # 0 ~ 1 scaling

        # up-sampling
        out = F.interpolate(out.unsqueeze(0).unsqueeze(0), size=[width, height], mode='bilinear')  # 4D로 바꿈

        batch_masks.append(out)
    batch_masks = torch.stack(batch_masks, dim=0).squeeze(1)  # [B, 1, 1, 240, 240] --> [B, 1, 240, 240]
    return batch_masks
//  [python]
def forward(self, x):
        input = x
        x = self.conv1(x)
        x = self.bn0(x)
        x = self.relu(x)
        convA_x = x
        x = self.maxpool(x)
        pooled_x = x

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

        gc_fm1 = self.br1(self.gcn1(fm1))
        gc_fm2 = self.br2(self.gcn2(fm2))
        gc_fm3 = self.br3(self.gcn3(fm3))
        gc_fm4 = self.br4(self.gcn4(fm4))

        gc_fm4 = F.upsample_bilinear(gc_fm4, fm3.size()[2:])
        gc_fm3 = F.upsample_bilinear(self.br5(gc_fm3 + gc_fm4), fm2.size()[2:])
        gc_fm2 = F.upsample_bilinear(self.br6(gc_fm2 + gc_fm3), fm1.size()[2:])
        gc_fm1 = F.upsample_bilinear(self.br7(gc_fm1 + gc_fm2), pooled_x.size()[2:])

        gc_fm1 = F.upsample_bilinear(self.br8(gc_fm1), scale_factor=2)

        out = F.upsample_bilinear(self.br9(gc_fm1), input.size()[2:])

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

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

        gc_fm1 = self.br1(self.gcn1(fm1))
        gc_fm2 = self.br2(self.gcn2(fm2))
        gc_fm3 = self.br3(self.gcn3(fm3))
        gc_fm4 = self.br4(self.gcn4(fm4))

        gc_fm4 = F.interpolate(gc_fm4, size=fm3.size()[2:], mode='bilinear')
        gc_fm3 = F.interpolate(self.br5(gc_fm3 + gc_fm4), size=fm2.size()[2:], mode='bilinear')
        gc_fm2 = F.interpolate(self.br6(gc_fm2 + gc_fm3), size=fm1.size()[2:], mode='bilinear')
        gc_fm1 = F.interpolate(self.br7(gc_fm1 + gc_fm2), size=pooled_x.size()[2:], mode='bilinear')

        gc_fm1 = F.interpolate(self.br8(gc_fm1), scale_factor=2, mode='bilinear')

        out = F.interpolate(self.br9(gc_fm1), size=input.size()[2:], mode='bilinear')

        return out
//  [python]
def forward(self, x, y, mask):
        while x.size()[3] > 1024:
            x, y = self.downsample(x), self.downsample(y)
        x_vgg, y_vgg = self.vgg(x), self.vgg(y)
        loss = 0
        for i in range(len(x_vgg)):
            _, ch, h, w = x_vgg[i].size()
            curr_mask = F.upsample_nearest(mask, (h, w))
            curr_mask = curr_mask.expand(-1, ch, -1, -1)
            loss += self.weights[i] * self.criterion(x_vgg[i] * curr_mask, y_vgg[i].detach() * curr_mask)
        return loss
//  [python]
def forward(self, x, y, mask):
        while x.size()[3] > 1024:
            x, y = self.downsample(x), self.downsample(y)
        x_vgg, y_vgg = self.vgg(x), self.vgg(y)
        loss = 0
        for i in range(len(x_vgg)):
            _, ch, h, w = x_vgg[i].size()
            curr_mask = F.interpolate(mask, (h, w), mode='nearest')
            curr_mask = curr_mask.expand(-1, ch, -1, -1)
            loss += self.weights[i] * self.criterion(x_vgg[i] * curr_mask, y_vgg[i].detach() * curr_mask)
        return loss
//  [python]
def forward(self, x):
        next_x = []
        for level, feature in enumerate(x):

            conv_args = dict()
            if self.offset is not None:
                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[level - 1], **conv_args))
            if level < len(x) - 1:
                temp_fea.append(F.upsample_bilinear(self.DyConv[0](x[level + 1], **conv_args),
                                                    size=[feature.size(2), feature.size(3)]))
            mean_fea = torch.mean(torch.stack(temp_fea), dim=0, keepdim=False)

            if self.AttnConv is not None:
                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.append(self.relu(mean_fea))

        return next_x
//  [python]
def forward(self, x):
        next_x = []
        for level, feature in enumerate(x):

            conv_args = dict()
            if self.offset is not None:
                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[level - 1], **conv_args))
            if level < len(x) - 1:
                temp_fea.append(F.interpolate(self.DyConv[0](x[level + 1], **conv_args),
                                                    size=[feature.size(2), feature.size(3)], mode='bilinear'))
            mean_fea = torch.mean(torch.stack(temp_fea), dim=0, keepdim=False)

            if self.AttnConv is not None:
                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.append(self.relu(mean_fea))

        return next_x
//  [python]
def forward(self, x_in):
        # pdb.set_trace()
        x_size = x_in.shape[2]
        y_size = x_in.shape[3]
        z_size = x_in.shape[4]

        hidden_states_out = self.swinViT(x_in, self.normalize)
        enc0 = self.encoder1(x_in)
        enc1 = self.encoder2(hidden_states_out[0])
        enc2 = self.encoder3(hidden_states_out[1])
        enc3 = self.encoder4(hidden_states_out[2])
        enc4 = self.encoder10(hidden_states_out[4])
        # pdb.set_trace()
        dec4 = F.relu(F.upsample(self.decoder5(enc4), size=(int(x_size / 8), int(y_size / 8), int(z_size / 8))))
        dec4 = dec4 + enc3
        dec3 = F.relu(F.upsample(self.decoder4(dec4), size=(int(x_size / 4), int(y_size / 4), int(z_size / 4))))
        dec3 = dec3 + enc2
        dec2 = F.relu(F.upsample(self.decoder3(dec3), size=(int(x_size / 2), int(y_size / 2), int(z_size / 2))))
        dec2 = dec2 + enc1
        dec1 = F.relu(F.upsample(self.decoder2(dec2), size=(x_size, y_size, z_size)))
        out = dec1 + enc0
        # pdb.set_trace()
        logits = self.out(out)
        return logits
//  [python]
def forward(self, x_in):
        x_size = x_in.shape[2]
        y_size = x_in.shape[3]
        z_size = x_in.shape[4]

        hidden_states_out = self.swinViT(x_in, self.normalize)
        enc0 = self.encoder1(x_in)
        enc1 = self.encoder2(hidden_states_out[0])
        enc2 = self.encoder3(hidden_states_out[1])
        enc3 = self.encoder4(hidden_states_out[2])
        enc4 = self.encoder10(hidden_states_out[4])
        dec4 = F.relu(F.interpolate(self.decoder5(enc4), size=(int(x_size / 8), int(y_size / 8), int(z_size / 8))))
        dec4 = dec4 + enc3
        dec3 = F.relu(F.interpolate(self.decoder4(dec4), size=(int(x_size / 4), int(y_size / 4), int(z_size / 4))))
        dec3 = dec3 + enc2
        dec2 = F.relu(F.interpolate(self.decoder3(dec3), size=(int(x_size / 2), int(y_size / 2), int(z_size / 2))))
        dec2 = dec2 + enc1
        dec1 = F.relu(F.interpolate(self.decoder2(dec2), size=(x_size, y_size, z_size)))
        out = dec1 + enc0
        logits = self.out(out)
        return logits
//  [python]
def batch_augment(self, images, attention_map, mode='crop', theta=0.5, padding_ratio=0.1):
        batches, _, imgH, imgW = images.size()

        if mode == 'crop':
            crop_images = []
            for batch_index in range(batches):
                atten_map = attention_map[batch_index:batch_index + 1]
                if isinstance(theta, tuple):
                    theta_c = random.uniform(*theta) * atten_map.max()
                else:
                    theta_c = theta * atten_map.max()

                crop_mask = F.upsample_bilinear(atten_map, size=(imgH, imgW)) >= theta_c
                nonzero_indices = torch.nonzero(crop_mask[0, 0, ...])
                height_min = max(int(nonzero_indices[:, 0].min().item() - padding_ratio * imgH), 0)
                height_max = min(int(nonzero_indices[:, 0].max().item() + padding_ratio * imgH), imgH)
                width_min = max(int(nonzero_indices[:, 1].min().item() - padding_ratio * imgW), 0)
                width_max = min(int(nonzero_indices[:, 1].max().item() + padding_ratio * imgW), imgW)

                crop_images.append(
                    F.upsample_bilinear(
                        images[batch_index:batch_index + 1, :, height_min:height_max, width_min:width_max],
                        size=(imgH, imgW)))
            crop_images = torch.cat(crop_images, dim=0)
            return crop_images

        elif mode == 'drop':
            drop_masks = []
            for batch_index in range(batches):
                atten_map = attention_map[batch_index:batch_index + 1]
                if isinstance(theta, tuple):
                    theta_d = random.uniform(*theta) * atten_map.max()
                else:
                    theta_d = theta * atten_map.max()

                drop_masks.append(F.upsample_bilinear(atten_map, size=(imgH, imgW)) < theta_d)
            drop_masks = torch.cat(drop_masks, dim=0)
            drop_images = images * drop_masks.float()
            return drop_images

        else:
            raise ValueError(
                'Expected mode in [\'crop\', \'drop\'], but received unsupported augmentation method %s' % mode)
//  [python]
def batch_augment(self, images, attention_map, mode='crop', theta=0.5, padding_ratio=0.1):
        batches, _, imgH, imgW = images.size()

        if mode == 'crop':
            crop_images = []
            for batch_index in range(batches):
                atten_map = attention_map[batch_index:batch_index + 1]
                if isinstance(theta, tuple):
                    theta_c = random.uniform(*theta) * atten_map.max()
                else:
                    theta_c = theta * atten_map.max()

                crop_mask = F.interpolate(atten_map, size=(imgH, imgW), mode='bilinear') >= theta_c
                nonzero_indices = torch.nonzero(crop_mask[0, 0, ...])
                height_min = max(int(nonzero_indices[:, 0].min().item() - padding_ratio * imgH), 0)
                height_max = min(int(nonzero_indices[:, 0].max().item() + padding_ratio * imgH), imgH)
                width_min = max(int(nonzero_indices[:, 1].min().item() - padding_ratio * imgW), 0)
                width_max = min(int(nonzero_indices[:, 1].max().item() + padding_ratio * imgW), imgW)

                crop_images.append(
                    F.interpolate(
                        images[batch_index:batch_index + 1, :, height_min:height_max, width_min:width_max],
                        size=(imgH, imgW), mode='bilinear'))
            crop_images = torch.cat(crop_images, dim=0)
            return crop_images

        elif mode == 'drop':
            drop_masks = []
            for batch_index in range(batches):
                atten_map = attention_map[batch_index:batch_index + 1]
                if isinstance(theta, tuple):
                    theta_d = random.uniform(*theta) * atten_map.max()
                else:
                    theta_d = theta * atten_map.max()

                drop_masks.append(F.interpolate(atten_map, size=(imgH, imgW), mode='bilinear') < theta_d)
            drop_masks = torch.cat(drop_masks, dim=0)
            drop_images = images * drop_masks.float()
            return drop_images

        else:
            raise ValueError(
                'Expected mode in [\'crop\', \'drop\'], but received unsupported augmentation method %s' % mode)
//  [python]
def forward(self, inp):
        """Get Inception feature maps

        Parameters
        ----------
        inp : torch.autograd.Variable
            Input tensor of shape Bx3xHxW. Values are expected to be in
            range (0, 1)

        Returns
        -------
        List of torch.autograd.Variable, corresponding to the selected output
        block, sorted ascending by index
        """
        outp = []
        x = inp

        if self.resize_input:
            x = F.upsample(x, size=(299, 299), mode='bilinear', align_corners=True)

        if self.normalize_input:
            x = x.clone()
            x[:, 0] = x[:, 0] * (0.229 / 0.5) + (0.485 - 0.5) / 0.5
            x[:, 1] = x[:, 1] * (0.224 / 0.5) + (0.456 - 0.5) / 0.5
            x[:, 2] = x[:, 2] * (0.225 / 0.5) + (0.406 - 0.5) / 0.5

        for idx, block in enumerate(self.blocks):
            x = block(x)
            if idx in self.output_blocks:
                outp.append(x)

            if idx == self.last_needed_block:
                break

        return outp
//  [python]
def forward(self, inp):
        """Get Inception feature maps

        Parameters
        ----------
        inp : torch.autograd.Variable
            Input tensor of shape Bx3xHxW. Values are expected to be in
            range (0, 1)

        Returns
        -------
        List of torch.autograd.Variable, corresponding to the selected output
        block, sorted ascending by index
        """
        outp = []
        x = inp

        if self.resize_input:
            x = F.interpolate(x, size=(299, 299), mode='bilinear', align_corners=True)

        if self.normalize_input:
            x = x.clone()
            x[:, 0] = x[:, 0] * (0.229 / 0.5) + (0.485 - 0.5) / 0.5
            x[:, 1] = x[:, 1] * (0.224 / 0.5) + (0.456 - 0.5) / 0.5
            x[:, 2] = x[:, 2] * (0.225 / 0.5) + (0.406 - 0.5) / 0.5

        for idx, block in enumerate(self.blocks):
            x = block(x)
            if idx in self.output_blocks:
                outp.append(x)

            if idx == self.last_needed_block:
                break

        return outp
//  [python]
def scale_tensor(input,size=512,mode='bilinear'):
    print(input.size())
    # b,h,w = input.size()
    _, _, h, w = input.size()
    if mode == 'nearest':
        if h == 512 and w == 512:
            return input
        return F.upsample_nearest(input,size=(size,size))
    if h>512 and w > 512:
        return F.upsample(input, size=(size,size), mode=mode, align_corners=True)
    return F.upsample(input, size=(size,size), mode=mode, align_corners=True)
//  [python]
def scale_tensor(input,size=512,mode='bilinear'):
    print(input.size())
    # b,h,w = input.size()
    _, _, h, w = input.size()
    if mode == 'nearest':
        if h == 512 and w == 512:
            return input
        return F.interpolate(input,size=(size,size), mode=mode)
    if h>512 and w > 512:
        return F.interpolate(input, size=(size,size), mode=mode, align_corners=True)
    return F.interpolate(input, size=(size,size), mode=mode, align_corners=True)