Instruction: The output of the PyTourch DL network doesn't match the last layer provided in the network

Input:
I am trying to build a PyTorch model that can predict the rank of a text where the output is a float number between 0 and 1.
My input details are

My batch size is 32.
Max length for the tokenizer is 116
In addition to the masks and Ids generated from the tokenizer, I am adding 11 values that were generated through preprocessing to the input text.

S the entire input shape would be 32 for batch and array with 127 item for each sample text provided
My layers are as follows:

a DistilBERT uncased transformer. and I am using the DistilBERT tokenizer over the text.
The following layer is a CNN that takes the output of the DistilBERT (127 channel) as input and provide 64 channels as output, with kernel=1
After this, 6 CNN layers each input is 64 and output is 64 with a kernel size of 3 and dilation increasing from 2 to 32. On top of each CNN, there is a relu and a maxpooling with 2 as kernal size.
My last CNN layer (and where the issue is happening) have 64 input channels and 32 output channels with a kernel size of 1 and a relu with AdaptiveMaxPool1d with size of 32 on top of it
Linear layer takes 32 and output 16
Linear layer takes 16 and output 1

below is my code
class Dataset(Dataset):
    def __init__(self, df, max_len, bert_model_name, multi=1):
        super().__init__()
        self.df = df.reset_index(drop=True)
        self.max_len = max_len

        self.tokenizer = DistilBertTokenizer.from_pretrained(
            bert_model_name,
            do_lower_case=True,
            strip_accents=True,
            wordpieces_prefix=None,
            use_fast=True
        )
        self.multiplier = multi

    def __getitem__(self, index):
        row = self.df.iloc[index]

        inputs = self.tokenizer.encode_plus(
            row.source,
            None,
            add_special_tokens=True,
            max_length=self.max_len,
            padding="max_length",
            return_token_type_ids=True,
            truncation=True
        )

        return (
            t.LongTensor(t.cat([
                t.LongTensor([
                    row.n_total_cells * self.multiplier,
                    row.n_code_cells * self.multiplier,
                    row.n_markdown_cells * self.multiplier,
                    row.word_counts * self.multiplier,
                    row.line_counts * self.multiplier,
                    row.empty_line_counts * self.multiplier,
                    row.full_lines_count * self.multiplier,
                    row.text_lines_count * self.multiplier,
                    row.tag_lines_count * self.multiplier,
                    row.weight * self.multiplier,
                    row.weight_counts * self.multiplier,
                ]),
                t.LongTensor(inputs['input_ids']),
            ], 0)),

            t.LongTensor(t.cat([
                t.ones(11, dtype=t.long),
                t.LongTensor(inputs['attention_mask']),
            ], 0)),
        )

class BModel(nn.Module):
    def __init__(self, bert_model_name):
        super(BModel, self).__init__()
        self.distill_bert = DistilBertModel.from_pretrained(bert_model_name)

        self.hidden_size = self.distill_bert.config.hidden_size
        print(self.hidden_size) # 768
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.3)

        self.cnn_layers()

    def forward(self, inputs):
        dbert = self.cnn_forward(inputs[0], inputs[1])
        return dbert

 def cnn_layers(self):
        self.layers = 4
        kernel_size = 3
        inp = 127
        out = 32
        grades = [2, 4, 8, 16, 32, 64, ]

        self.convs = nn.ModuleList()
        self.relus = nn.ModuleList()
        self.maxs = nn.ModuleList()
        self.norms = nn.ModuleList()

        self.start_conv = nn.Conv1d(
            in_channels=inp,
            out_channels=64,
            kernel_size=1,
            bias=True
        )

        for i in range(self.layers):
            # dilated convolutions
            self.convs.append(nn.Conv1d(
                in_channels=64,
                out_channels=64,
                kernel_size = kernel_size,
                bias=False,
                dilation=grades[i]
            ))

            self.relus.append(nn.ReLU())

            self.maxs.append(nn.MaxPool1d(
                kernel_size=kernel_size-1,
            ))

            self.norms.append(nn.BatchNorm1d(
                num_features=64,
            ))

        self.end_conv = nn.Conv1d(
            in_channels=64,
            out_channels=out,
            kernel_size=1,
            bias=True
        )

        self.max_pool = nn.AdaptiveMaxPool1d(out)

        self.top1 = nn.Linear(out, 16)
        self.top2 = nn.Linear(16, 1)

    def cnn_forward(self, ids, masks):
        x = self.distill_bert(ids, masks)[0]
        x = self.relu(x)
        x = self.dropout(x)
        print(f"X size after BERT:", x.size())

        x = self.start_conv(x)
        print(f"X size after First Conv:", x.size())
        for i in range(self.layers):
            x = self.norms[i](self.maxs[i](self.relus[i](self.convs[i](x))))
            print(f"X size after {i} CNN dilation:", x.size())

        x = self.max_pool(t.abs(self.end_conv(x)))
        print("X size after AdaptiveMaxPool1d:", x.size())

        x = self.top1(x)
        print("X size after before-last linear:", x.size())

        x = self.top2(x)
        print("X size after last linear:", x.size())
        return x

Printing the output size after each layer would be as below
X size after First Conv: torch.Size([32, 64, 768])
X size after 0 CNN dilation: torch.Size([32, 64, 382])
X size after 1 CNN dilation: torch.Size([32, 64, 187])
X size after 2 CNN dilation: torch.Size([32, 64, 85])
X size after 3 CNN dilation: torch.Size([32, 64, 26])
X size after AdaptiveMaxPool1d: torch.Size([32, 32, 32])
X size after before-last linear: torch.Size([32, 32, 16])
X size after last linear: torch.Size([32, 32, 1]

The issue I am facing is after the AdaptiveMaxPool1d, the output of this layer suppose to be 2 dimensions instead of 3 [32, 32] instead of [32, 32, 32]
The output of AdaptiveMaxPool1d fits into the linear layer but is with one extra dimension causing the output pred to differ from the true input
when I check the pred size vs the true size it would be
y_pred shape (12480,)
y_val shape (390,)

and the code blow with this error
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Input In [13], in <cell line: 21>()
     17 # print(mkdn_train_loader, mkdn_val_loader)
     18
     19 ########################################################################################################################

File E:\KAGGLE_COMP\pt_model.py:796, in train(model, train_loader, val_loader, epochs, patience, path)
    793 print('y_val shape', y_val.shape)
    794 print(y_pred[:10])
--> 796 print("Validation MSE:", np.round(mean_squared_error(y_val, y_pred), 4))
    797 print()
    799 early_stopping(np.round(mean_squared_error(y_val, y_pred), 4), model)

File ~\AppData\Local\Programs\Python\Python310\lib\site-packages\sklearn\metrics\_regression.py:438, in mean_squared_error(y_true, y_pred, sample_weight, multioutput, squared)
    378 def mean_squared_error(
    379     y_true, y_pred, *, sample_weight=None, multioutput="uniform_average", squared=True
    380 ):
    381     """Mean squared error regression loss.
    382
    383     Read more in the :ref:`User Guide <mean_squared_error>`.
   (...)
    436     0.825...
    437     """
--> 438     y_type, y_true, y_pred, multioutput = _check_reg_targets(
    439         y_true, y_pred, multioutput
    440     )
    441     check_consistent_length(y_true, y_pred, sample_weight)
    442     output_errors = np.average((y_true - y_pred) ** 2, axis=0, weights=sample_weight)

File ~\AppData\Local\Programs\Python\Python310\lib\site-packages\sklearn\metrics\_regression.py:94, in _check_reg_targets(y_true, y_pred, multioutput, dtype)
     60 def _check_reg_targets(y_true, y_pred, multioutput, dtype="numeric"):
     61     """Check that y_true and y_pred belong to the same regression task.
     62
     63     Parameters
   (...)
     92         the dtype argument passed to check_array.
     93     """
---> 94     check_consistent_length(y_true, y_pred)
     95     y_true = check_array(y_true, ensure_2d=False, dtype=dtype)
     96     y_pred = check_array(y_pred, ensure_2d=False, dtype=dtype)

File ~\AppData\Local\Programs\Python\Python310\lib\site-packages\sklearn\utils\validation.py:332, in check_consistent_length(*arrays)
    330 uniques = np.unique(lengths)
    331 if len(uniques) > 1:
--> 332     raise ValueError(
    333         "Found input variables with inconsistent numbers of samples: %r"
    334         % [int(l) for l in lengths]
    335     )

ValueError: Found input variables with inconsistent numbers of samples: [390, 12480]

I need to know what I must change to make this run and the size is passed with correct shape.

Output:
From AdaptiveMaxPool1d Documentation:
If the Input is in Shape of (N, C, L_in), then your output would be in the shape of (N, C, L_out).
Since your input shape to the AdaptiveMaxPool1d is in the shape of (32, 32, 26) and you've set the output_size to 32 ( value of "out" variable ), your output shape comes out as (32, 32, 32).
I suggest to set output_size as 1 and use squeeze(2) to squish down the dimension.
Something like this:
# For initialization of maxpool layer.
nn.AdaptiveMaxPool1d(1)
# ---------
# In forward add squeeze(2) after max_pool like this:
x = self.max_pool(t.abs(self.end_conv(x))).squeeze(2)


Instruction: Equivalent to tokenizer() in Transformers 2.5.0?

Input:
I am trying to convert the following code to work with Transformers 2.5.0. As written, it works in version 4.18.0, but not 2.5.0.
# Converting pretrained BERT classification model to regression model
# i.e. extracting base model and swapping out heads

from transformers import BertTokenizer, BertModel, BertConfig, BertForMaskedLM, BertForSequenceClassification, AutoConfig, AutoModelForTokenClassification
import torch
import numpy as np

old_model = BertForSequenceClassification.from_pretrained("textattack/bert-base-uncased-yelp-polarity")
model = BertForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=1)
model.bert = old_model.bert

# Ensure that model parameters are equivalent except for classifier head layer
for param_name in model.state_dict():
    if 'classifier' not in param_name:
        sub_param, full_param = model.state_dict()[param_name], old_model.state_dict()[param_name] # type: torch.Tensor, torch.Tensor
        assert (sub_param.cpu().numpy() == full_param.cpu().numpy()).all(), param_name

tokenizer = BertTokenizer.from_pretrained("textattack/bert-base-uncased-yelp-polarity")
inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")

with torch.no_grad():
    logits = model(**inputs).logits

output_value = np.array(logits)[0][0]
print(output_value)

tokenizer is not callable with transformers 2.5.0, resulting the following:
TypeError                                 Traceback (most recent call last)
<ipython-input-1-d83f0d613f4b> in <module>
     19
     20
---> 21 inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
     22
     23 with torch.no_grad():

TypeError: 'BertTokenizer' object is not callable

However, attempting to replace tokenizer() with tokenizer.tokenize() results in the following:
TypeError                                 Traceback (most recent call last)
<ipython-input-2-1d431131eb87> in <module>
     21
     22 with torch.no_grad():
---> 23     logits = model(**inputs).logits
     24
     25 output_value = np.array(logits)[0][0]

TypeError: BertForSequenceClassification object argument after ** must be a mapping, not list

Any help would be greatly appreciated.

Solution
Using tokenizer.encode_plus() as suggested by @cronoik:
tokenized = tokenizer.encode_plus("Hello, my dog is cute", return_tensors="pt")

with torch.no_grad():
    logits = model(**tokenized)

output_value = np.array(logits)[0]
print(output_value)

Output:
Sadly their documentation for the old versions is broken, but you can use encode_plus as shown in the following (he oldest available documentation of encode_plus is from 2.10.0):
import torch
from transformers import BertTokenizer

t = BertTokenizer.from_pretrained("textattack/bert-base-uncased-yelp-polarity")
tokenized = t.encode_plus("Hello, my dog is cute", return_tensors='pt')
print(tokenized)

Output:
{'input_ids': tensor([[  101,  7592,  1010,  2026,  3899,  2003, 10140,   102]]),
'token_type_ids': tensor([[0, 0, 0, 0, 0, 0, 0, 0]]),
'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1]])}


Instruction: When is `stage is None` in pytorch lightning?

Input:
Some official pytorch lightning docs have code that refer to stage as Optional[str] with for example the following code
import pytorch_lightning as pl
from torch.utils.data import random_split, DataLoader

# Note - you must have torchvision installed for this example
from torchvision.datasets import MNIST
from torchvision import transforms

class MNISTDataModule(pl.LightningDataModule):
    def __init__(self, data_dir: str = "./"):
        super().__init__()
        self.data_dir = data_dir
        self.transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])

    def prepare_data(self):
        # download
        MNIST(self.data_dir, train=True, download=True)
        MNIST(self.data_dir, train=False, download=True)

    def setup(self, stage: Optional[str] = None):

        # Assign train/val datasets for use in dataloaders
        if stage == "fit" or stage is None:
            mnist_full = MNIST(self.data_dir, train=True, transform=self.transform)
            self.mnist_train, self.mnist_val = random_split(mnist_full, [55000, 5000])

        # Assign test dataset for use in dataloader(s)
        if stage == "test" or stage is None:
            self.mnist_test = MNIST(self.data_dir, train=False, transform=self.transform)

        if stage == "predict" or stage is None:
            self.mnist_predict = MNIST(self.data_dir, train=False, transform=self.transform)

    def train_dataloader(self):
        return DataLoader(self.mnist_train, batch_size=32)

    def val_dataloader(self):
        return DataLoader(self.mnist_val, batch_size=32)

    def test_dataloader(self):
        return DataLoader(self.mnist_test, batch_size=32)

    def predict_dataloader(self):
        return DataLoader(self.mnist_predict, batch_size=32)

When does stage take the value of None? I could find no docs describing this.

Output:
The Trainer will never send stage=None to the setup hook, or any of the other hooks that take this argument. The type is annotated optional and the default value is None for historical reasons. The values will always be one of "fit", "validate", "test", "predict".
There is an RFC to change this to a required argument to avoid confusion. The link provides some more context why it has been like this for the past.


Instruction: How can I stop pytorch model from downloading vgg .pth pretrained file every time I do inference?

Input:
I'm trying to use the u-net segmentation model at https://github.com/khanhha/crack_segmentation, and incorporate it into my pipeline. However, I noticed that whenever I use 'inference_unet.py', for the first time in the session, it downloads a .pth file for vgg.
Downloading: "https://download.pytorch.org/models/vgg16-397923af.pth" to C:\Users\hedey/.cache\torch\hub\checkpoints\vgg16-397923af.pth

It's not practical to download that file every time I make an inference, especially that this will be a part of an application. How can I avoid having to download that file every time?
Here is the code at 'inference_unet.py':
import sys
import os
import numpy as np
from pathlib import Path
import cv2 as cv
import torch
import torch.nn.functional as F
from torch.autograd import Variable
import torchvision.transforms as transforms
from unet.unet_transfer import UNet16, input_size
import matplotlib.pyplot as plt
import argparse
from os.path import join
from PIL import Image
import gc
from utils import load_unet_vgg16, load_unet_resnet_101, load_unet_resnet_34
from tqdm import tqdm

def evaluate_img(model, img):
    input_width, input_height = input_size[0], input_size[1]

    img_1 = cv.resize(img, (input_width, input_height), cv.INTER_AREA)
    X = train_tfms(Image.fromarray(img_1))
    X = Variable(X.unsqueeze(0)).cuda()  # [N, 1, H, W]

    mask = model(X)

    mask = F.sigmoid(mask[0, 0]).data.cpu().numpy()
    mask = cv.resize(mask, (img_width, img_height), cv.INTER_AREA)
    return mask

def evaluate_img_patch(model, img):
    input_width, input_height = input_size[0], input_size[1]

    img_height, img_width, img_channels = img.shape

    if img_width < input_width or img_height < input_height:
        return evaluate_img(model, img)

    stride_ratio = 0.1
    stride = int(input_width * stride_ratio)

    normalization_map = np.zeros((img_height, img_width), dtype=np.int16)

    patches = []
    patch_locs = []
    for y in range(0, img_height - input_height + 1, stride):
        for x in range(0, img_width - input_width + 1, stride):
            segment = img[y:y + input_height, x:x + input_width]
            normalization_map[y:y + input_height, x:x + input_width] += 1
            patches.append(segment)
            patch_locs.append((x, y))

    patches = np.array(patches)
    if len(patch_locs) <= 0:
        return None

    preds = []
    for i, patch in enumerate(patches):
        patch_n = train_tfms(Image.fromarray(patch))
        X = Variable(patch_n.unsqueeze(0)).cuda()  # [N, 1, H, W]
        masks_pred = model(X)
        mask = F.sigmoid(masks_pred[0, 0]).data.cpu().numpy()
        preds.append(mask)

    probability_map = np.zeros((img_height, img_width), dtype=float)
    for i, response in enumerate(preds):
        coords = patch_locs[i]
        probability_map[coords[1]:coords[1] + input_height, coords[0]:coords[0] + input_width] += response

    return probability_map

def disable_axis():
    plt.axis('off')
    plt.gca().axes.get_xaxis().set_visible(False)
    plt.gca().axes.get_yaxis().set_visible(False)
    plt.gca().axes.get_xaxis().set_ticklabels([])
    plt.gca().axes.get_yaxis().set_ticklabels([])

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-img_dir',type=str, help='input dataset directory')
    parser.add_argument('-model_path', type=str, help='trained model path')
    parser.add_argument('-model_type', type=str, choices=['vgg16', 'resnet101', 'resnet34'])
    parser.add_argument('-out_viz_dir', type=str, default='', required=False, help='visualization output dir')
    parser.add_argument('-out_pred_dir', type=str, default='', required=False,  help='prediction output dir')
    parser.add_argument('-threshold', type=float, default=0.2 , help='threshold to cut off crack response')
    args = parser.parse_args()

    if args.out_viz_dir != '':
        os.makedirs(args.out_viz_dir, exist_ok=True)
        for path in Path(args.out_viz_dir).glob('*.*'):
            os.remove(str(path))

    if args.out_pred_dir != '':
        os.makedirs(args.out_pred_dir, exist_ok=True)
        for path in Path(args.out_pred_dir).glob('*.*'):
            os.remove(str(path))

    if args.model_type == 'vgg16':
        model = load_unet_vgg16(args.model_path)
    elif args.model_type  == 'resnet101':
        model = load_unet_resnet_101(args.model_path)
    elif args.model_type  == 'resnet34':
        model = load_unet_resnet_34(args.model_path)
        print(model)
    else:
        print('undefind model name pattern')
        exit()

    channel_means = [0.485, 0.456, 0.406]
    channel_stds  = [0.229, 0.224, 0.225]

    paths = [path for path in Path(args.img_dir).glob('*.*')]
    for path in tqdm(paths):
        #print(str(path))

        #train_tfms = transforms.Compose([transforms.ToTensor(), transforms.Normalize(channel_means, channel_stds)])
        train_tfms = transforms.Compose([transforms.ToTensor()])

        img_0 = Image.open(str(path))
        img_0 = np.asarray(img_0)

        if len(img_0.shape) != 3:
            print(f'incorrect image shape: {path.name}{img_0.shape}')
            continue

        img_0 = img_0[:,:,:3]
        img_height, img_width, img_channels = img_0.shape

        #img_height, img_width = img_0.shape

        prob_map_full = evaluate_img(model, img_0)

        if args.out_pred_dir != '':
            #cv.imwrite(filename=join(args.out_pred_dir, f'{path.stem}.jpg'), img=(prob_map_full * 255).astype(np.uint8))
            cv.imwrite(filename=join(args.out_pred_dir, f'{path.stem}.jpg'), img=(prob_map_full).astype(np.uint8))

        if args.out_viz_dir != '':
            # plt.subplot(121)
            # plt.imshow(img_0), plt.title(f'{img_0.shape}')
            if img_0.shape[0] > 2000 or img_0.shape[1] > 2000:
                img_1 = cv.resize(img_0, None, fx=0.2, fy=0.2, interpolation=cv.INTER_AREA)
            else:
                img_1 = img_0

            # plt.subplot(122)
            # plt.imshow(img_0), plt.title(f'{img_0.shape}')
            # plt.show()

            prob_map_patch = evaluate_img_patch(model, img_1)

            #plt.title(f'name={path.stem}. \n cut-off threshold = {args.threshold}', fontsize=4)
            prob_map_viz_patch = prob_map_patch.copy()
            prob_map_viz_patch = prob_map_viz_patch/ prob_map_viz_patch.max()
            prob_map_viz_patch[prob_map_viz_patch < args.threshold] = 0.0
            fig = plt.figure()
            st = fig.suptitle(f'name={path.stem} \n cut-off threshold = {args.threshold}', fontsize="x-large")
            ax = fig.add_subplot(231)
            ax.imshow(img_1)
            ax = fig.add_subplot(232)
            ax.imshow(prob_map_viz_patch)
            ax = fig.add_subplot(233)
            ax.imshow(img_1)
            ax.imshow(prob_map_viz_patch, alpha=0.4)

            prob_map_viz_full = prob_map_full.copy()
            prob_map_viz_full[prob_map_viz_full < args.threshold] = 0.0

            ax = fig.add_subplot(234)
            ax.imshow(img_0)
            ax = fig.add_subplot(235)
            ax.imshow(prob_map_viz_full)
            ax = fig.add_subplot(236)
            ax.imshow(img_0)
            ax.imshow(prob_map_viz_full, alpha=0.4)

            plt.savefig(join(args.out_viz_dir, f'{path.stem}.jpg'), dpi=500)
            plt.close('all')

        gc.collect()

Here is the code at 'utils.py':
import json
from datetime import datetime
from pathlib import Path

import random
import numpy as np

import torch
import tqdm
from unet.unet_transfer import UNet16, UNetResNet

class AverageMeter(object):
    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count

def cuda(x):
    #return x.cuda(async=True) if torch.cuda.is_available() else x
    return x.cuda(non_blocking=True) if torch.cuda.is_available() else x

def write_event(log, step, **data):
    data['step'] = step
    data['dt'] = datetime.now().isoformat()
    log.write(json.dumps(data, sort_keys=True))
    log.write('\n')
    log.flush()

def check_crop_size(image_height, image_width):
    """Checks if image size divisible by 32.
    Args:
        image_height:
        image_width:
    Returns:
        True if both height and width divisible by 32 and False otherwise.
    """
    return image_height % 32 == 0 and image_width % 32 == 0

def create_model(device, type ='vgg16'):
    assert type == 'vgg16' or type == 'resnet101'
    if type == 'vgg16':
        model = UNet16(pretrained=True)
    elif type == 'resnet101':
        model = UNetResNet(pretrained=True, encoder_depth=101, num_classes=1)
    else:
        assert False
    model.eval()
    return model.to(device)

def load_unet_vgg16(model_path):
    model = UNet16(pretrained=True)
    #model = UNet16(pretrained=False)
    checkpoint = torch.load(model_path)
    if 'model' in checkpoint:
        model.load_state_dict(checkpoint['model'])
    elif 'state_dict' in checkpoint:
        model.load_state_dict(checkpoint['check_point'])
    else:
        raise Exception('undefind model format')

    model.cuda()
    model.eval()

    return model

def load_unet_resnet_101(model_path):
    #model = UNetResNet(pretrained=True, encoder_depth=101, num_classes=1)
    model = UNetResNet(pretrained=True, encoder_depth=101, num_classes=8)
    checkpoint = torch.load(model_path)
    if 'model' in checkpoint:
        model.load_state_dict(checkpoint['model'])
    elif 'state_dict' in checkpoint:
        model.load_state_dict(checkpoint['check_point'])
    else:
        raise Exception('undefind model format')

    model.cuda()
    model.eval()

    return model

def load_unet_resnet_34(model_path):
    model = UNetResNet(pretrained=True, encoder_depth=34, num_classes=1)
    checkpoint = torch.load(model_path)
    if 'model' in checkpoint:
        model.load_state_dict(checkpoint['model'])
    elif 'state_dict' in checkpoint:
        model.load_state_dict(checkpoint['check_point'])
    else:
        raise Exception('undefind model format')

    model.cuda()
    model.eval()

    return model

def train(args, model, criterion, train_loader, valid_loader, validation, init_optimizer, n_epochs=None, fold=None,
          num_classes=None):
    lr = args.lr
    n_epochs = n_epochs or args.n_epochs
    optimizer = init_optimizer(lr)

    root = Path(args.model_path)
    model_path = root / 'model_{fold}.pt'.format(fold=fold)
    if model_path.exists():
        state = torch.load(str(model_path))
        epoch = state['epoch']
        step = state['step']
        model.load_state_dict(state['model'])
        print('Restored model, epoch {}, step {:,}'.format(epoch, step))
    else:
        epoch = 1
        step = 0

    save = lambda ep: torch.save({
        'model': model.state_dict(),
        'epoch': ep,
        'step': step,
    }, str(model_path))

    report_each = 10
    log = root.joinpath('train_{fold}.log'.format(fold=fold)).open('at', encoding='utf8')
    valid_losses = []
    for epoch in range(epoch, n_epochs + 1):
        model.train()
        random.seed()
        tq = tqdm.tqdm(total=(len(train_loader) * args.batch_size))
        tq.set_description('Epoch {}, lr {}'.format(epoch, lr))
        losses = []
        tl = train_loader
        try:
            mean_loss = 0
            for i, (inputs, targets) in enumerate(tl):
                inputs = cuda(inputs)

                with torch.no_grad():
                    targets = cuda(targets)

                outputs = model(inputs)
                #print(outputs.shape, targets.shape)
                loss = criterion(outputs, targets)
                optimizer.zero_grad()
                batch_size = inputs.size(0)
                loss.backward()
                optimizer.step()
                step += 1
                tq.update(batch_size)
                losses.append(loss.item())
                mean_loss = np.mean(losses[-report_each:])
                tq.set_postfix(loss='{:.5f}'.format(mean_loss))
                if i and i % report_each == 0:
                    write_event(log, step, loss=mean_loss)
            write_event(log, step, loss=mean_loss)
            tq.close()
            save(epoch + 1)
            valid_metrics = validation(model, criterion, valid_loader, num_classes)
            write_event(log, step, **valid_metrics)
            valid_loss = valid_metrics['valid_loss']
            valid_losses.append(valid_loss)
        except KeyboardInterrupt:
            tq.close()
            print('Ctrl+C, saving snapshot')
            save(epoch)
            print('done.')
            return

Here is the code at 'unet_transfer.py':
from torch import nn
from torch.nn import functional as F
import torch
from torchvision import models
import torchvision

input_size = (448, 448)

class Interpolate(nn.Module):
    def __init__(self, size=None, scale_factor=None, mode='nearest', align_corners=False):
        super(Interpolate, self).__init__()
        self.interp = nn.functional.interpolate
        self.size = size
        self.mode = mode
        self.scale_factor = scale_factor
        self.align_corners = align_corners

    def forward(self, x):
        x = self.interp(x, size=self.size, scale_factor=self.scale_factor,
                        mode=self.mode, align_corners=self.align_corners)
        return x

def conv3x3(in_, out):
    return nn.Conv2d(in_, out, 3, padding=1)

class ConvRelu(nn.Module):
    def __init__(self, in_, out):
        super().__init__()
        self.conv = conv3x3(in_, out)
        self.activation = nn.ReLU(inplace=True)

    def forward(self, x):
        x = self.conv(x)
        x = self.activation(x)
        return x

class DecoderBlockV2(nn.Module):
    def __init__(self, in_channels, middle_channels, out_channels, is_deconv=True):
        super(DecoderBlockV2, self).__init__()
        self.in_channels = in_channels

        if is_deconv:
            """
                Paramaters for Deconvolution were chosen to avoid artifacts, following
                link https://distill.pub/2016/deconv-checkerboard/
            """

            #self.block = nn.ModuleList(
            self.block = nn.Sequential(
                ConvRelu(in_channels, middle_channels),
                nn.ConvTranspose2d(middle_channels, out_channels, kernel_size=4, stride=2,
                                   padding=1),
                nn.ReLU(inplace=True)
            )
        else:
            self.block = nn.Sequential(
                Interpolate(scale_factor=2, mode='bilinear'),
                ConvRelu(in_channels, middle_channels),
                ConvRelu(middle_channels, out_channels),
            )

    def forward(self, x):
        return self.block(x)

class UNet16(nn.Module):
    def __init__(self, num_classes=1, num_filters=32, pretrained=False, is_deconv=False):
    #def __init__(self, num_classes=8, num_filters=32, pretrained=False, is_deconv=False):
        """
        :param num_classes:
        :param num_filters:
        :param pretrained:
            False - no pre-trained network used
            True - encoder pre-trained with VGG16
        :is_deconv:
            False: bilinear interpolation is used in decoder
            True: deconvolution is used in decoder
        """
        super().__init__()
        self.num_classes = num_classes

        self.pool = nn.MaxPool2d(2, 2)

        #print(torchvision.models.vgg16(pretrained=pretrained))

        self.encoder = torchvision.models.vgg16(pretrained=pretrained).features
        #self.encoder = torchvision.models.vgg16(pretrained=False).features

        self.relu = nn.ReLU(inplace=True)

        self.conv1 = nn.Sequential(self.encoder[0],
                                   self.relu,
                                   self.encoder[2],
                                   self.relu)

        self.conv2 = nn.Sequential(self.encoder[5],
                                   self.relu,
                                   self.encoder[7],
                                   self.relu)

        self.conv3 = nn.Sequential(self.encoder[10],
                                   self.relu,
                                   self.encoder[12],
                                   self.relu,
                                   self.encoder[14],
                                   self.relu)

        self.conv4 = nn.Sequential(self.encoder[17],
                                   self.relu,
                                   self.encoder[19],
                                   self.relu,
                                   self.encoder[21],
                                   self.relu)

        self.conv5 = nn.Sequential(self.encoder[24],
                                   self.relu,
                                   self.encoder[26],
                                   self.relu,
                                   self.encoder[28],
                                   self.relu)

        self.center = DecoderBlockV2(512, num_filters * 8 * 2, num_filters * 8, is_deconv)

        self.dec5 = DecoderBlockV2(512 + num_filters * 8, num_filters * 8 * 2, num_filters * 8, is_deconv)
        self.dec4 = DecoderBlockV2(512 + num_filters * 8, num_filters * 8 * 2, num_filters * 8, is_deconv)
        self.dec3 = DecoderBlockV2(256 + num_filters * 8, num_filters * 4 * 2, num_filters * 2, is_deconv)
        self.dec2 = DecoderBlockV2(128 + num_filters * 2, num_filters * 2 * 2, num_filters, is_deconv)
        self.dec1 = ConvRelu(64 + num_filters, num_filters)
        self.final = nn.Conv2d(num_filters, num_classes, kernel_size=1)

    def forward(self, x):
        conv1 = self.conv1(x)
        conv2 = self.conv2(self.pool(conv1))
        conv3 = self.conv3(self.pool(conv2))
        conv4 = self.conv4(self.pool(conv3))
        conv5 = self.conv5(self.pool(conv4))

        center = self.center(self.pool(conv5))

        dec5 = self.dec5(torch.cat([center, conv5], 1))

        dec4 = self.dec4(torch.cat([dec5, conv4], 1))
        dec3 = self.dec3(torch.cat([dec4, conv3], 1))
        dec2 = self.dec2(torch.cat([dec3, conv2], 1))
        dec1 = self.dec1(torch.cat([dec2, conv1], 1))

        if self.num_classes > 1:
            x_out = F.log_softmax(self.final(dec1), dim=1)
        else:
            x_out = self.final(dec1)
            #x_out = F.sigmoid(x_out)

        return x_out

class UNetResNet(nn.Module):

    def __init__(self, encoder_depth, num_classes, num_filters=32, dropout_2d=0.2,
                 pretrained=False, is_deconv=False):
        super().__init__()
        self.num_classes = num_classes
        self.dropout_2d = dropout_2d

        if encoder_depth == 34:
            self.encoder = torchvision.models.resnet34(pretrained=pretrained)
            bottom_channel_nr = 512
        elif encoder_depth == 101:
            self.encoder = torchvision.models.resnet101(pretrained=pretrained)
            bottom_channel_nr = 2048
        elif encoder_depth == 152:
            self.encoder = torchvision.models.resnet152(pretrained=pretrained)
            bottom_channel_nr = 2048
        else:
            raise NotImplementedError('only 34, 101, 152 version of Resnet are implemented')

        self.pool = nn.MaxPool2d(2, 2)

        self.relu = nn.ReLU(inplace=True)

        #self.conv1 = nn.Sequential(self.encoder.conv1,
        #                           self.encoder.bn1,
        #                           self.encoder.relu,
        #                           self.pool)

        self.conv1 = nn.Sequential(nn.Conv2d(1,64,kernel_size=(7,7),stride=(2,2),padding=(3,3),bias=False), # 1 Here is for grayscale images, replace by 3 if you need RGB/BGR
                                   nn.BatchNorm2d(64),
                                   nn.ReLU(),
                                   self.pool
                                )

        self.conv2 = self.encoder.layer1

        self.conv3 = self.encoder.layer2

        self.conv4 = self.encoder.layer3

        self.conv5 = self.encoder.layer4

        self.center = DecoderBlockV2(bottom_channel_nr, num_filters * 8 * 2, num_filters * 8, is_deconv)
        self.dec5 = DecoderBlockV2(bottom_channel_nr + num_filters * 8, num_filters * 8 * 2, num_filters * 8, is_deconv)
        self.dec4 = DecoderBlockV2(bottom_channel_nr // 2 + num_filters * 8, num_filters * 8 * 2, num_filters * 8,
                                   is_deconv)
        self.dec3 = DecoderBlockV2(bottom_channel_nr // 4 + num_filters * 8, num_filters * 4 * 2, num_filters * 2,
                                   is_deconv)
        self.dec2 = DecoderBlockV2(bottom_channel_nr // 8 + num_filters * 2, num_filters * 2 * 2, num_filters * 2 * 2,
                                   is_deconv)
        self.dec1 = DecoderBlockV2(num_filters * 2 * 2, num_filters * 2 * 2, num_filters, is_deconv)
        self.dec0 = ConvRelu(num_filters, num_filters)
        self.final = nn.Conv2d(num_filters, num_classes, kernel_size=1)
        #self.final = nn.Conv2d(num_filters, 1, kernel_size=1)

    def forward(self, x):
        conv1 = self.conv1(x)
        conv2 = self.conv2(conv1)
        conv3 = self.conv3(conv2)
        conv4 = self.conv4(conv3)
        conv5 = self.conv5(conv4)

        pool = self.pool(conv5)
        center = self.center(pool)

        dec5 = self.dec5(torch.cat([center, conv5], 1))

        dec4 = self.dec4(torch.cat([dec5, conv4], 1))
        dec3 = self.dec3(torch.cat([dec4, conv3], 1))
        dec2 = self.dec2(torch.cat([dec3, conv2], 1))
        dec1 = self.dec1(dec2)
        dec0 = self.dec0(dec1)

        return self.final(F.dropout2d(dec0, p=self.dropout_2d))

'''
class UNetResNet(nn.Module):
    """PyTorch U-Net model using ResNet(34, 101 or 152) encoder.
    UNet: https://arxiv.org/abs/1505.04597
    ResNet: https://arxiv.org/abs/1512.03385
    Proposed by Alexander Buslaev: https://www.linkedin.com/in/al-buslaev/
    Args:
            encoder_depth (int): Depth of a ResNet encoder (34, 101 or 152).
            num_classes (int): Number of output classes.
            num_filters (int, optional): Number of filters in the last layer of decoder. Defaults to 32.
            dropout_2d (float, optional): Probability factor of dropout layer before output layer. Defaults to 0.2.
            pretrained (bool, optional):
                False - no pre-trained weights are being used.
                True  - ResNet encoder is pre-trained on ImageNet.
                Defaults to False.
            is_deconv (bool, optional):
                False: bilinear interpolation is used in decoder.
                True: deconvolution is used in decoder.
                Defaults to False.
    """

    def __init__(self, encoder_depth, num_classes, num_filters=32, dropout_2d=0.2,
                 pretrained=False, is_deconv=False):
        super().__init__()
        self.num_classes = num_classes
        self.dropout_2d = dropout_2d

        if encoder_depth == 34:
            self.encoder = torchvision.models.resnet34(pretrained=pretrained)
            bottom_channel_nr = 512
        elif encoder_depth == 101:
            self.encoder = torchvision.models.resnet101(pretrained=pretrained)
            bottom_channel_nr = 2048
        elif encoder_depth == 152:
            self.encoder = torchvision.models.resnet152(pretrained=pretrained)
            bottom_channel_nr = 2048
        else:
            raise NotImplementedError('only 34, 101, 152 version of Resnet are implemented')

        self.pool = nn.MaxPool2d(2, 2)

        self.relu = nn.ReLU(inplace=True)

        self.conv1 = nn.Sequential(self.encoder.conv1,
                                   self.encoder.bn1,
                                   self.encoder.relu,
                                   self.pool)

        self.conv2 = self.encoder.layer1

        self.conv3 = self.encoder.layer2

        self.conv4 = self.encoder.layer3

        self.conv5 = self.encoder.layer4

        self.center = DecoderBlockV2(bottom_channel_nr, num_filters * 8 * 2, num_filters * 8, is_deconv)
        self.dec5 = DecoderBlockV2(bottom_channel_nr + num_filters * 8, num_filters * 8 * 2, num_filters * 8, is_deconv)
        self.dec4 = DecoderBlockV2(bottom_channel_nr // 2 + num_filters * 8, num_filters * 8 * 2, num_filters * 8,
                                   is_deconv)
        self.dec3 = DecoderBlockV2(bottom_channel_nr // 4 + num_filters * 8, num_filters * 4 * 2, num_filters * 2,
                                   is_deconv)
        self.dec2 = DecoderBlockV2(bottom_channel_nr // 8 + num_filters * 2, num_filters * 2 * 2, num_filters * 2 * 2,
                                   is_deconv)
        self.dec1 = DecoderBlockV2(num_filters * 2 * 2, num_filters * 2 * 2, num_filters, is_deconv)
        self.dec0 = ConvRelu(num_filters, num_filters)
        self.final = nn.Conv2d(num_filters, num_classes, kernel_size=1)

    def forward(self, x):
        self.conv1 = torch.nn.Conv2d(1, 64, (7, 7), (2, 2), (3, 3), bias=False)
        conv1 = self.conv1(x)
        conv2 = self.conv2(conv1)
        conv3 = self.conv3(conv2)
        conv4 = self.conv4(conv3)
        conv5 = self.conv5(conv4)

        pool = self.pool(conv5)
        center = self.center(pool)

        dec5 = self.dec5(torch.cat([center, conv5], 1))

        dec4 = self.dec4(torch.cat([dec5, conv4], 1))
        dec3 = self.dec3(torch.cat([dec4, conv3], 1))
        dec2 = self.dec2(torch.cat([dec3, conv2], 1))
        dec1 = self.dec1(dec2)
        dec0 = self.dec0(dec1)

        return self.final(F.dropout2d(dec0, p=self.dropout_2d))
'''

Output:
In 'utils.py', 'load_unet_vgg16' function receives a path to a checkpoint to load the model from, so the initialization of the weights step (which happens before loading from the checkpoint) is unnecessary.
The function 'load_unet_vgg16' could be as follows:
def load_unet_vgg16(model_path, pretrained=True):
    model = UNet16(pretrained=pretrained)
    checkpoint = torch.load(model_path)
    if 'model' in checkpoint:
        model.load_state_dict(checkpoint['model'])
    elif 'state_dict' in checkpoint:
        model.load_state_dict(checkpoint['check_point'])
    else:
        raise Exception('undefind model format')

    model.cuda()
    model.eval()

    return model

Then you pass a path to a checkpoint and 'pretrained=False' since you are making an inference.


Instruction: Masking dimensions outside of a window

Input:
I have a Pytorch tensor t of shape (n, x, y), and I'd like to apply a mask such that, for all y > x + k (with k being a constant), t[n, x, y] = -inf.
I believe I can do this with advanced indexing, but can't figure out how.
If not, a simple way to do this is to construct a mask like that with loops (slow, but do it one time and cache it), and then t += mask, since -inf + z == -inf for all z.
Is there a better way to do this?

Output:
Notice that the condition y ≥ x corresponds to the upper triangle, while y > x is the strict upper triangle. Therefore y > x + k is the upper triangle part with a shift equal to 1 + k.
You can construct a triangle mask using torch.triu which actually allows for a shift argument named diagonal, referring to the position of the diagonal. Assign the desired value, here -torch.inf, using this mask will allow you to obtain the desired result.
Overall it comes down to:
>>> m = torch.ones_like(t, dtype=bool).triu(1+k)
>>> t[m] = -torch.inf

Alternatively, a one-liner is possible using torch.where:
>>> torch.where(torch.ones_like(t).bool().triu(1+k), -torch.inf, t)

Since the mask is equal for all batch elements, you can get away with creating a single 2D mask and masking t on its 2nd and 3rd axes:
>>> m = torch.ones_like(t[0], dtype=bool).triu(1+k)
>>> t[:,m] = -torch.inf


Instruction: PyTorch Lightning - Display metrics after validation epoch

Input:
I've implemented validation_epoch_end to produce and log metrics, and when I run trainer.validate, the metrics appear in my notebook.
However, when I run trainer.fit, only the training metrics appear; not the validation ones.
The validation step is still being run (because the validation code calls a print statement, which does appear), but the validation metrics don't appear, even though they're logged.  Or, if they do appear, the next epoch immediately erases them, so that I can't see them.
(Likewise, tensorboard sees the validation metrics)
How can I see the validation epoch end metrics in a notebook, as each epoch occurs?

Output:
You could do the following. Let's say you have the following LightningModule:
class MNISTModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.l1 = torch.nn.Linear(28 * 28, 10)

    def forward(self, x):
        return torch.relu(self.l1(x.view(x.size(0), -1)))

    def training_step(self, batch, batch_nb):
        x, y = batch
        loss = F.cross_entropy(self(x), y)
        # prog_bar=True will display the value on the progress bar statically for the last complete train epoch
        self.log("train_loss", loss, on_step=False, on_epoch=True, prog_bar=True)

        return loss

    def validation_step(self, batch, batch_nb):
        x, y = batch
        loss = F.cross_entropy(self(x), y)
        # prog_bar=True will display the value on the progress bar statically for the last complete validation epoch
        self.log("val_loss", loss, on_step=False, on_epoch=True, prog_bar=True)

        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.02)

The trick is to use prog_bar=True in combination with on_step and on_epoch depending on when you want the update on the progress bar. So, in this case, when training:
# Train the model ⚡
trainer.fit(mnist_model, MNIST_dm)

you will see:
Epoch 4: 100% -------------------------- 939/939 [00:09<00:00, 94.51it/s, loss=0.636, v_num=4, val_loss=0.743, train_loss=0.726]

Where loss will be updating each batch as it is the step loss. However, val_loss and train_loss will be static values that will only change after each validation or train epoch respectively.


Instruction: How to compute partial derivatives of a component of a vector-valued function?

Input:
Let’s say I have a function Psi with a 4-dimensional vector output, that takes a 3-dimensional vector u as input. I would like to compute the gradient of the first three components of Psi w.r.t. the respective three components of u:
import torch

u = torch.tensor([1.,2.,3.], requires_grad=True)

psi = torch.zeros(4)
psi[0] = 2*u[0]
psi[1] = 2*u[1]
psi[2] = 2*u[2]
psi[3] = torch.dot(u,u)

grad_Psi_0 = torch.autograd.grad(psi[0], u[0])
grad_Psi_1 = torch.autograd.grad(psi[1], u[1])
grad_Psi_2 = torch.autograd.grad(psi[2], u[2])

And I get the error that u[0],u[1], and u[2] are not used in the graph:
---> 19 grad_Psi_0 = torch.autograd.grad(psi[0], u[0])
     20 grad_Psi_1 = torch.autograd.grad(psi[1], u[1])
     21 grad_Psi_2 = torch.autograd.grad(psi[2], u[2])

File ~/.local/lib/python3.10/site-packages/torch/autograd/__init__.py:275, in grad(outputs, inputs, grad_outputs, retain_graph, create_graph, only_inputs, allow_unused, is_grads_batched)
    273     return _vmap_internals._vmap(vjp, 0, 0, allow_none_pass_through=True)(grad_outputs)
    274 else:
--> 275     return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
    276         outputs, grad_outputs_, retain_graph, create_graph, inputs,
    277         allow_unused, accumulate_grad=False)

RuntimeError: One of the differentiated Tensors appears to not have been used in the graph. Set allow_unused=True if this is the desired behavior.

This is strange, as components of u are used to compute components of psi, so it should be possible to compute derivatives of the components of psi w.r.t to components of u. How to fix this?
Answer: based on answer from @Ivan, for higher-order derivatives, component-wise calls to gradient require create_graph=True, otherwise the same error happens as described above.
import torch

# u = grad(Phi) = [2*u0, 2*u1, 2*u2]
# Phi = u0**2 + u1**2 + u2**2 = dot(u,u)

u = torch.tensor([1.,2.,3.], requires_grad=True)

psi = torch.zeros(4)
psi[0] = 2*u[0]
psi[1] = 2*u[1]
psi[2] = 2*u[2]
psi[3] = torch.dot(u,u)

print("u = ",u)
print("psi = ",psi)

grad_v_x = torch.autograd.grad(psi[0], u, retain_graph=True)[0]
print(grad_v_x)
grad_v_y = torch.autograd.grad(psi[1], u, retain_graph=True)[0]
print(grad_v_y)
grad_v_z = torch.autograd.grad(psi[2], u, retain_graph=True)[0]
print(grad_v_z)
div_v = grad_v_x[0] + grad_v_y[1] + grad_v_z[2]

# Divergence of the vector phi[0:3]=2u0 + 2u1 + 2u2 w.r.t [u0,u1,u2] = 2+2+2=6
print (div_v)

# laplace(psi[3]) = \partial_u0^2 psi[3] + \partial_u1^2 psi[3] + \partial_u2^2 psi[3]
# = \partial_u0 2x + \partial_u1 2u1 + \partial_u2 2u2 = 2 + 2 + 2 = 6
d_phi_du = torch.autograd.grad(psi[3], u, create_graph=True, retain_graph=True)[0]
print(d_phi_du)
dd_phi_d2u0 = torch.autograd.grad(d_phi_du[0], u, retain_graph=True)[0]
dd_phi_d2u1 = torch.autograd.grad(d_phi_du[1], u, retain_graph=True)[0]
dd_phi_d2u2 = torch.autograd.grad(d_phi_du[2], u, retain_graph=True)[0]

laplace_phi = torch.dot(dd_phi_d2u0 + dd_phi_d2u1 + dd_phi_d2u2, torch.ones(3))

print(laplace_phi)

Output:
The reason why is because u[0] is actually a copy so the one used on the following line:
psi[0] = 2*u[0]

is different to the one used here
grad_Psi_0 = torch.autograd.grad(psi[0], u[0])

which means they are not linked in the computation graph.
A possible solution is to assign u[0] to a separate variable such that it can be used on both calls:
>>> u0 = u[0]
>>> psi[0] = 2*u0
>>> torch.autograd.grad(psi[0], u0)
(tensor(2.),)

Alternatively, you can call autograd.grad directly on u:
>>> psi[0] = 2*u[0]
>>> torch.autograd.grad(psi[0], u)
(tensor([2., 0., 0.]),)


Instruction: How to fix the input dimension from convolution flatten to feed forward layer?

Input:
I am using nni framework on python to do Neural Architecture Search. In that I have defined model as:
from nni.nas.pytorch import mutables

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = mutables.LayerChoice([
            nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1),
            nn.Conv2d(3, 32, kernel_size=5, stride=1, padding=1)
        ])  # try 3x3 kernel and 5x5 kernel
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout2d(0.25)
        self.dropout2 = nn.Dropout2d(0.5)
        self.fc1 = nn.Linear(14400, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x) #Here is error coming
        x = F.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        output = F.log_softmax(x, dim=1)
        return output

What the above code does apart from building the model is it also gives a choice to below algorithm to choose between two layers as the first convolution layer, either layer with 3X3 kernel or 5X5 kernel.
Also I am new to pyTorch so let me know if you can already see a mistake in above.
Moving on, it is coupled by below code:
dataset_train = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=True, transform=transform)
dataset_valid = torchvision.datasets.CIFAR10(root='./data', train=False,
                                       download=True, transform=transform)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), 0.05, momentum=0.9, weight_decay=1.0E-4)

# use NAS here
def top1_accuracy(output, target):
    # this is the function that computes the reward, as required by ENAS algorithm
    batch_size = target.size(0)
    _, predicted = torch.max(output.data, 1)
    return (predicted == target).sum().item() / batch_size

def metrics_fn(output, target):
    # metrics function receives output and target and computes a dict of metrics
    return {"acc1": top1_accuracy(output, target)}

from nni.algorithms.nas.pytorch import enas
trainer = enas.EnasTrainer(model,
                           loss=criterion,
                           metrics=metrics_fn,
                           reward_function=top1_accuracy,
                           optimizer=optimizer,
                           batch_size=128,
                           num_epochs=10,  # 10 epochs
                           dataset_train=dataset_train,
                           dataset_valid=dataset_valid,
                           log_frequency=10)  # print log every 10 steps
trainer.train()  # training
trainer.export(file="model_dir/final_architecture.json")  # export the final architecture to file

What the above does is downloads and gets cifar10 dataset, uses the above generated model to train on it and finds which model performs best (based on two choices of layers, you can have more choices as well). But it raises an error:
     22         x = self.dropout1(x)
     23         x = torch.flatten(x, 1)
---> 24         x = self.fc1(x)
     25         x = F.relu(x)
     26         x = self.dropout2(x)

/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
   1128         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1129                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1130             return forward_call(*input, **kwargs)
   1131         # Do not call functions when jit is used
   1132         full_backward_hooks, non_full_backward_hooks = [], []

/usr/local/lib/python3.7/dist-packages/torch/nn/modules/linear.py in forward(self, input)
    112
    113     def forward(self, input: Tensor) -> Tensor:
--> 114         return F.linear(input, self.weight, self.bias)
    115
    116     def extra_repr(self) -> str:

RuntimeError: mat1 and mat2 shapes cannot be multiplied (128x12544 and 14400x128)

I know this is because the flatten layer converts it to a dimension which is not what the first fully connected layer expects. When I do convert it to what the error says, I get the below error:
RuntimeError: mat1 and mat2 shapes cannot be multiplied (128x14400 and 12544x128)

I believe it happens because of the choice in first convolution layer. My question is how do I fix this? And if nni or something feels not understandable to you, there is the option of just putting the dimensions of fully connected layer as number of hidden units in that layer without mentioning  the input in KERAS. But I suppose pyTorch requires input dimension to be correctly put, is there a way I can just say after flatten, to go for a hidden fully connected layer with just the number of units and not the input shape as well which I believe is causing the problems?

Output:
For conv with kernel_zise=5 you need to padding=2 and not 1.
Fix:
        self.conv1 = mutables.LayerChoice([
            nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1),
            nn.Conv2d(3, 32, kernel_size=5, stride=1, padding=1)
        ])

to
        self.conv1 = mutables.LayerChoice([
            nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1),
            nn.Conv2d(3, 32, kernel_size=5, stride=1, padding=2)  # match padding size to kernel size
        ])

Update:
Recent versions of pytorch allow you to specify padding='same' and avoid the need to come up with the correct value for padding.
However, I strongly urge you to use the formula for computing the output shape of a convolution layer (found here) and manually compute the correct value for padding. This is a good sanity check to ensure you understand what you are doing.


Instruction: How to train a network with two datasets and two output heads?

Input:
I am tying to train two datasets at the same time to get good results on both datasets.
data_loader_iterator = iter(data_loader_second)
for batch_idx, (image1, label1) in enumerate(data_loader):
    image1 = image1.to(args.local_rank)
    label1 = label1.to(args.local_rank)
    label1 = label1.squeeze()
    try:
        image2, label2 = next(data_loader_iterator)
    except StopIteration:
        data_loader_iterator = iter(data_loader_second)
        image2, label2 = next(data_loader_iterator)
    image2 = image2.to(args.local_rank)
    label2 = label2.to(args.local_rank)
    label2 = label2.squeeze()

    embedding1 = backbone.forward(image1)
    embedding2 = backbone.forward(image2)
    output1 = head1.forward(embedding1, label1)
    output2 = head2.forward(embedding2, label2)
    loss1 = criterion(output1, label1)
    loss2 = criterion(output2, label2)
    loss = loss1 + loss2

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    ...

head1 = HeadFactory(args.head_type, args.head_conf_file, 751577).get_head()
head2 = HeadFactory(args.head_type, args.head_conf_file, 253788).get_head()
...
optimizer = torch.optim.AdamW(params=[{"params": backbone.parameters()},{"params": head1.parameters()},{"params": head2.parameters()}], lr = args.lr, weight_decay=0.05)
...
criterion = torch.nn.CrossEntropyLoss().to(args.local_rank)

This can't work correctly WARNING: torch.distributed.elastic.multiprocessing.api: Sending process 2919 closing signal SIGTERM and I wonder how to declare the optimizer

Output:
It's unclear from your error what exactly went wrong and where.
Let's try a few things:

How do you merge two datasets with different labels into the same DataLoader? how can you ensure that the first label belongs to head1 while the second to head2?
I would recommend two Datasets with two DataLoaders to be iterated using a zip function:

for batch_idx, (image1, label1), (image2, label2) in enumerate(zip(data_loader1, dataloader2)):
  # ...

I wonder if the way you pass all the parameters to the optimizer is not causing trouble.
How about wrapping everything into a single nn.Container?

container = nn.ModuleList([backbone, head1, head2])
optimizer = torch.optim.AdamW(container.parameters(), lr=...)

If you need to wrap the model for parallelism, you can wrap the container.

Additionally, Are you training in a multi-GPU settings? In that case, you need to add a few pieces to the code:

Wrap both backbone and heads with a distributed module.
You data loaders need to be aware of parallelism. This is usually done via DistributedSampler.


Instruction: How to use `self` as a function, like `self(**data)` in the example code

Input:
Here is part of code from mmcv (mmcv/runner/base_module.py).
I don't understand how does self(**data) works.

class BaseDetector(BaseModule, metaclass=ABCMeta):
    """Base class for detectors."""

    # ...

    def train_step(self, data, optimizer):

        losses = self(**data)
        loss, log_vars = self._parse_losses(losses)

        outputs = dict(
            loss=loss, log_vars=log_vars, num_samples=len(data['img_metas']))

        return outputs

class BaseModule(nn.Module, metaclass=ABCMeta):
    def __init__(self, init_cfg: Optional[dict] = None):
    # ...

Tried to find some function like __call__ in the parent class, but it seems not the right direction to go.
Could any one explain why self could be used as a function?

Output:
This is nothing specific to PyTorch but rather to how Python works.
"Calling" an object is equivalent to calling a special Python function named __call__. Essentially, this is the function which implements the call behaviour. In your case self(**data) will call BaseDetector's __call__ function which has most probably been implemented by one of the parent classes.
Here is a minimal example to understand this behaviour:
class A():
    def __call__(self):
        return 'hello'

class B(A):
    def foo(self):
        return self() + ' world'

Initialize a new object from B:
b = B()

Then let us call the foo method on b:
>>> b.foo()
'hello world'

This function has been calling b() i.e. b.__call__():
>>> b() # same as b.__call__()
'hello'

As an additional note, you can explicitly call the implementation of the parent class. This is useful when you need to overload the original implementation in your child's class.
It is done by calling super:
class A():
    def __call__(self):
        return 'hello'

class B(A):
    def __call__(self):
        # calling self() here would create an infinite recursion loop
        return super().__call__() + ' world'

Then you can call an instance of B the same way:
>>> b = B()
>>> b()
'hello world'


Instruction: Plot Pytorch vectors with TSNE

Input:
I am using the ESM-1b  model to train it with some protein sequences. I already have the vectors and now I wanted to plot them using TSNE. However, when I try to pass the vectors to the TSNE model I get:
'list' object has no attribute 'shape'`

How should I plot the Pytorch vectors (they are Pytorch tensors, actually)?
The code I have so far:
sequence_representations = []
for i, (_, seq) in enumerate(new_list):
   sequence_representations.append(token_representations[i, 1 : len(seq) + 1].mean(0))

This is an example of the Pytorch tensors I have (sequence_representations):
[tensor([-0.0054,  0.1090, -0.0046,  ...,  0.0465,  0.0426, -0.0675]),
 tensor([-0.0025,  0.0228, -0.0521,  ..., -0.0611,  0.1010, -0.0103]),
 tensor([ 0.1168, -0.0189, -0.0121,  ..., -0.0388,  0.0586, -0.0285]),......

TSNE:
X_embedded = TSNE(n_components=2, learning_rate='auto', init='random').fit_transform(sequence_representations) #Where I get the error

Output:
Assuming you are using scipy's TSNE, you'll need sequence_representations to be

ndarray of shape (n_samples, n_features)

Right now have a list of pytorch tensors.
To convert sequence_representations to a numpy ndarray you'll need:
seq_np = torch.stack(sequence_representations)  # from list of 1d tensors to a 2d tensor
seq_np = seq_np.numpy()  # convert to numpy


Instruction: Cannot load pretrained model for continue training

Input:
I trained my own model but decided to continue training.
When I use the code below, my model shows high BCELoss as it is a non-trained model.
Where is the problem?
Thank you
model_1 = SimpleCnn(n_classes) # model class
model_1.load_state_dict(torch.load('./model.pth', map_location='cuda:0'))
model_1.to(DEVICE) # torch cuda device
history = train(train_dataset, val_dataset, model=model_1, epochs=8, batch_size=16) # train function
torch.save(model_1.state_dict(), 'model_1.pth')

Output:
In order to continue training, you need to save not only the state_dict of the model, but of the optimizer's as well.
That is, during training, you need to save not only the trained weights of the model, but some other parameters. For example:
def train_function(...):
  for e in range(num_epochs):
    ...
  # when done training
  torch.save({'model', model.state_dict(),
              'opt', optimizer.state_dict(),
              'lr', lr_sched.state_dict(),  # you also need to save the state of the learning rate scheduler
              ...  # there might be other things that define the "state" of your training
             }, PATH)

Then, if you resume training:
checkpoint = torch.load(PATH)
model.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['opt'])
lr_sched.load_state_dict(checkpoint['lr'])
... # might need to restore other things


Instruction: No `predict_dataloader()` method defined to run `Trainer.predict`

Input:
I am trying to get predictions from my model based on the test set dataloader (will want to save both x and y^hat, which I need to test later on).
I tried:
my_results = trainer.predict(model = model, datamodule=dm)

With the following items present in my code:
class TimeseriesDataset(Dataset):
    '''
    Custom Dataset subclass.
    Serves as input to DataLoader to transform X
      into sequence data using rolling window.
    DataLoader using this dataset will output batches
      of `(batch_size, seq_len, n_features)` shape.
    Suitable as an input to RNNs.
    '''
    def __init__(self, X: np.ndarray, y: np.ndarray, seq_len: int = 1):
        self.X = torch.tensor(X).float()
        self.y = torch.tensor(y).float()
        self.seq_len = seq_len

    def __len__(self):
        return self.X.__len__() - (self.seq_len-1)

    def __getitem__(self, index):
        return (self.X[index:index+self.seq_len], self.y[index+self.seq_len-1])

class LSTMRegressor(pl.LightningModule):
    '''
    Standard PyTorch Lightning module:
    https://pytorch-lightning.readthedocs.io/en/latest/lightning_module.html
    '''
    def __init__(self,
                 n_features,
                 hidden_size,
                 seq_len,
                 batch_size,
                 num_layers,
                 dropout,
                 learning_rate,
                 criterion):
        super(LSTMRegressor, self).__init__()
        self.n_features = n_features
        self.hidden_size = hidden_size
        self.seq_len = seq_len
        self.batch_size = batch_size
        self.num_layers = num_layers
        self.dropout = dropout
        self.criterion = criterion
        self.learning_rate = learning_rate

        self.lstm = nn.LSTM(input_size=n_features,
                            hidden_size=hidden_size,
                            num_layers=num_layers,
                            dropout=dropout,
                            batch_first=True)
        self.linear = nn.Linear(hidden_size, 2)

    def forward(self, x):
        # lstm_out = (batch_size, seq_len, hidden_size)
        lstm_out, _ = self.lstm(x)
        y_pred = self.linear(lstm_out[:,-1])
        return y_pred

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=self.learning_rate)

    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        loss = self.criterion(y_hat, y)
        # result = pl.TrainResult(loss)
        self.log('train_loss', loss)
        return loss

def predict_step(self, batch, batch_idx):
    with torch.no_grad():
        x, y = batch
        y_hat = self(x)
        return x, y_hat
    def validation_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        loss = self.criterion(y_hat, y)
        # result = pl.EvalResult(checkpoint_on=loss)
        self.log('val_loss', loss)
        enable_checkpointing = True, #ModelCheckpoint(monitor='val_loss')
        # checkpoint_callback = ModelCheckpoint(
        # monitor='val_loss',
        # dirpath='./lstm',
        # filename='lstm{epoch:02d}-val_loss{val/loss:.2f}',
        # auto_insert_metric_name=False
    # )
        return loss

    def test_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        loss = self.criterion(y_hat, y)
        # result = pl.EvalResult()
        self.log('test_loss', loss)
        enable_checkpointing = True, #ModelCheckpoint(monitor='test_loss')   #TODO check if loss is the thing to return in this function
        return loss

and:
class CryptoDataModule(pl.LightningDataModule):
    '''
    PyTorch Lighting DataModule subclass:
    https://pytorch-lightning.readthedocs.io/en/latest/datamodules.html

    Serves the purpose of aggregating all data loading
      and processing work in one place.
    '''

    def __init__(self, seq_len = 1, batch_size = 128, num_workers=0):
        super().__init__()
        self.seq_len = seq_len
        self.batch_size = batch_size
        self.num_workers = num_workers
        self.X_train = None
        self.y_train = None
        self.X_val = None
        self.y_val = None
        self.X_test = None
        self.X_test = None
        self.columns = None
        self.preprocessing = None

def setup(self, stage=None):
    '''
    Data is resampled to hourly intervals.
    Both 'np.nan' and '?' are converted to 'np.nan'
    'Date' and 'Time' columns are merged into 'dt' index
    '''

    if stage == 'fit' and self.X_train is not None:
        return
    if stage == 'test' or stage == 'predict' and self.X_test is not None:
        return
    if stage is None and self.X_train is not None and self.X_test is not None:
        return

    path = './eth_data_1d.csv'

    df = pd.read_csv(
        path,
        sep=',',
        infer_datetime_format=True,
        low_memory=False,
        na_values=['nan','?'],
        index_col='Time'
    )

    y = pd.concat([df['Top'], df['Btm']], axis=1, keys=['Top', 'Btm'])

    X = df.dropna().copy()
    self.columns = X.columns

    X_cv, X_test, y_cv, y_test = train_test_split(
        X, y, test_size=0.2, shuffle=False
    )

    X_train, X_val, y_train, y_val = train_test_split(
        X_cv, y_cv, test_size=0.25, shuffle=False
    )

    preprocessing = StandardScaler()
    preprocessing.fit(X_train)

    self.X_train = preprocessing.transform(X_train)
    self.y_train = y_train.values.reshape((-1, 2))
    self.X_val = preprocessing.transform(X_val)
    self.y_val = y_val.values.reshape((-1, 2))
    self.X_test = preprocessing.transform(X_test)
    self.y_test = y_test.values.reshape((-1, 2))

def train_dataloader(self):
        train_dataset = TimeseriesDataset(self.X_train,
                                          self.y_train,
                                          seq_len=self.seq_len)
        train_loader = DataLoader(train_dataset,
                                  batch_size = self.batch_size,
                                  shuffle = False,
                                  num_workers = self.num_workers)

        return train_loader

    def val_dataloader(self):
        val_dataset = TimeseriesDataset(self.X_val,
                                        self.y_val,
                                        seq_len=self.seq_len)
        val_loader = DataLoader(val_dataset,
                                batch_size = self.batch_size,
                                shuffle = False,
                                num_workers = self.num_workers)

        return val_loader

    def test_dataloader(self):
        test_dataset = TimeseriesDataset(self.X_test,
                                         self.y_test,
                                         seq_len=self.seq_len)
        test_loader = DataLoader(test_dataset,
                                 batch_size = self.batch_size,
                                 shuffle = False,
                                 num_workers = self.num_workers)

        return test_loader

Gives me the following error:
MisconfigurationException                 Traceback (most recent call last)
/Users/xxx/ai_bt/model.ipynb Cell 22 in <cell line: 34>()
      1 # train on test set too! : see below
      2 # trainer.test(dataloaders=test_dataloaders)
      3
   (...)
     30 # with torch.no_grad():
     31     # predictions = trainer.predict(model, dm)
---> 34 my_results = trainer.predict(model = model, datamodule=dm)

File /opt/homebrew/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py:1025, in Trainer.predict(self, model, dataloaders, datamodule, return_predictions, ckpt_path)
   1000 r"""
   1001 Run inference on your data.
   1002 This will call the model forward function to compute predictions. Useful to perform distributed
   (...)
   1022     Returns a list of dictionaries, one for each provided dataloader containing their respective predictions.
   1023 """
   1024 self.strategy.model = model or self.lightning_module
-> 1025 return self._call_and_handle_interrupt(
   1026     self._predict_impl, model, dataloaders, datamodule, return_predictions, ckpt_path
   1027 )

File /opt/homebrew/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py:723, in Trainer._call_and_handle_interrupt(self, trainer_fn, *args, **kwargs)
    721         return self.strategy.launcher.launch(trainer_fn, *args, trainer=self, **kwargs)
...
--> 197     raise MisconfigurationException(f"No `{loader_name}()` method defined to run `Trainer.{trainer_method}`.")
    199 # predict_step is not required to be overridden
    200 if stage == "predict":

MisconfigurationException: No `predict_dataloader()` method defined to run `Trainer.predict`.

It must be something silly because I can't figure out what the dataloader refers to. the dm argument works fine during training...
Update: based on @Mikel B's answer I added:
def predict_dataloader(self):
        predict_dataset = TimeseriesDataset(self.X_test,
                                         self.y_test,
                                         seq_len=self.seq_len)
        predict_loader = DataLoader(predict_dataset,
                                 batch_size = self.batch_size,
                                 shuffle = False,
                                 num_workers = self.num_workers)

        return predict_loader

Which results in:
---> 44     output = model(batch)
     45     output = model.proba(output) # if not part of forward already
     46     prediction_list.append(output)

File /opt/homebrew/lib/python3.9/site-packages/torch/nn/modules/module.py:1130, in Module._call_impl(self, *input, **kwargs)
   1126 # If we don't have any hooks, we want to skip the rest of the logic in
   1127 # this function, and just call forward.
   1128 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1129         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1130     return forward_call(*input, **kwargs)
   1131 # Do not call functions when jit is used
   1132 full_backward_hooks, non_full_backward_hooks = [], []

/Users/user/ai_bt/model.ipynb Cell 22 in LSTMRegressor.forward(self, x)
     32 def forward(self, x):
     33     # lstm_out = (batch_size, seq_len, hidden_size)
---> 34     lstm_out, _ = self.lstm(x)
     35     y_pred = self.linear(lstm_out[:,-1])
     36     return y_pred
...
--> 731     is_batched = input.dim() == 3
    732     batch_dim = 0 if self.batch_first else 1
    733     if not is_batched:

AttributeError: 'list' object has no attribute 'dim'

Output:
You have not defined predict_dataloader() in you LightningDataModule:
class MNISTDataModule(pl.LightningDataModule):
    def __init__(self, data_dir: str = "path/to/dir", batch_size: int = 32):
        super().__init__()
        self.data_dir = data_dir
        self.batch_size = batch_size

    def setup(self, stage: Optional[str] = None):
        self.mnist_test = MNIST(self.data_dir, train=False)
        self.mnist_predict = MNIST(self.data_dir, train=False)
        mnist_full = MNIST(self.data_dir, train=True)
        self.mnist_train, self.mnist_val = random_split(mnist_full, [55000, 5000])

    def train_dataloader(self):
        return DataLoader(self.mnist_train, batch_size=self.batch_size)

    def val_dataloader(self):
        return DataLoader(self.mnist_val, batch_size=self.batch_size)

    def test_dataloader(self):
        return DataLoader(self.mnist_test, batch_size=self.batch_size)

    # THIS IS WHAT YOU ARE MISSING
    def predict_dataloader(self):
        return DataLoader(self.mnist_predict, batch_size=self.batch_size)

    def teardown(self, stage: Optional[str] = None):
        # Used to clean-up when the run is finished
        ...

Without this method the trainer does not know which data to load for the predict_step


Instruction: How to get the inference compute graph of the pytorch model?

Input:
I want to hand write a framework to perform inference of a given neural network. The network is so complicated, so to make sure my implementation is correct, I need to know how exactly the inference process is done on device.
I tried to use torchviz to visualize the network, but what I got seems to be the back propagation compute graph, which is really hard to understand.

Then I tried to convert the pytorch model to ONNX format, following the instruction enter link description here, but when I tried to visualize it, it seems that the original layers of the model had been seperated into very small operators.

I just want to get the result like this

How can I get this? Thanks!

Output: Have you tried saving the model with torch.save (https://pytorch.org/tutorials/beginner/saving_loading_models.html) and opening it with Netron? The last view you showed is a view of the Netron app.


Instruction: Label Smoothing in PyTorch - Using BCE loss -> doing it with the data itself

Input:
i am doing a classification task (binary) in PyTorch, so with labels 0 und 1.
No I want introduce label smoothing as another regularization technique.
Because I Use the ice loss, there is no such function to use label smoothing as
in the cross entropy loss (for man than 0,1).
Now I am considering to implement it not in the loss but in the data itself.
Would it be right to just replace my y_true to for example 0->0.1 and 1->0.9
before they go into the loss?

Output:
You can replace the 0 with 0.1  and 1 with 0.9 if label smoothing is 0.1
criterion(disc_fake_pred, torch.zeros_like(disc_fake_pred)+0.1) #0.1
criterion(disc_real_pred, torch.ones_like(disc_real_pred)-0.1)  #0.9


Instruction: Multi-channel, 2D mask weights using BCEWithLogitsLoss in Pytorch

Input:
I have a set of 256x256 images that are each labeled with nine, binary 256x256 masks. I am trying to calculate the pos_weight in order to weight the BCEWithLogitsLoss using Pytorch.
The shape of my masks tensor is tensor([1000, 9, 256, 256]) where 1000 is the number of training images, 9 is the number of mask channels (all encoded to 0/1), and 256 is the size of each image side.
To calculate pos_weight, I have summed the zeros in each mask, and divided that number by the sum of all of the ones in each mask (following the advice suggested here.):
(masks[:,channel,:,:]==0).sum()/masks[:,channel,:,:].sum()

Calculating the weight for every mask channel provides a tensor with the shape of tensor([9]), which seems intuitive to me, since I want a pos_weight value for each of the nine mask channels. However when I try to fit my model, I get the following error message:
RuntimeError: The size of tensor a (9) must match the size of
tensor b (256) at non-singleton dimension 3

This error message is surprising because it suggests that the weights need to be the size of one of the image sides, but not the number of mask channels. What shape should pos_weight be and how do I specify that it should be providing weights for the mask channels instead of the image pixels?

Output:
TLDR; This is a broadcasting issue which is surprisingly not handled by PyTorch's nn.BCEWithLogitsLoss namely F.binary_cross_entropy_with_logits. It might actually be worth putting out a Github issue linking to this SO thread to notify the developers of this undesirable behaviour.
In the documentation page of nn.BCEWithLogitsLoss, it is stated that the provided positive weights tensor pos_weight:

Must be a vector with length equal to the number of classes.

This is of course what you were expecting (rightly so) since positive weights refer to the weight given to the positive instances for every single class. Since your prediction and target tensors are multi-dimensional this seems to not be handled properly by PyTorch.

Anyhows, here is a minimal example showing how you can bypass this error and also showing the manual computation of the binary cross-entropy, as reference.
Here is the setup of the prediction and target tensors pred and label respectively:
>>> c=2;b=5;h=3;w=3
>>> pred = torch.rand(b,c,h,w)
>>> label = torch.randint(0,2, (b,c,h,w), dtype=float)

Now for the definition of the positive weight, notice the leading singletons dimensions:
>>> pos_weight = torch.rand(c,1,1)

In your case, with your existing 1D tensor of length c, you would simply have to unsqueeze two extra dimensions for the height and width dimensions. This means doing something like: pos_weight = pos_weight[:,None,None].
Calling the bce with logits function or its oop equivalent:
>>> F.binary_cross_entropy_with_logits(pred, label, pos_weight=pos_weight).mean()

Which is equivalent, in plain code to:
>>> z = torch.sigmoid(pred)
>>> bce = -(pos_weight*label*torch.log(z) + (1-label)*torch.log(1-z))

Note, that the built-in function would have the desired behaviour (i.e. no error message) if the class dimension was last in your prediction and target tensors.
>>> pos_weight = torch.rand(c)
>>> F.binary_cross_entropy_with_logits(
...    pred.transpose(1,-1),
...    label.transpose(1,-1),
...    pos_weight=pos_weight)

In other words, we are applying the function with format NHWC which means the pos_weight of format C can be multiplied properly. So the result above effectively yields the same result as:
>>> F.binary_cross_entropy_with_logits(
...    pred,
...    label,
...    pos_weight=pos_weight[:,None,None])

You can read more about the pos_weight in BCEWithLogitsLoss in another thread here


Instruction: How to index an array/tensor on the "rest" of the indices given some indices?

Input:
Given some array (or tensor):
x = np.array([[0, 1, 0, 0, 0],
              [0, 0, 0, 1, 0],
              [1, 0, 0, 0, 0]])

and some indices of dimension equaling the number of rows in x:
idx = np.array([3, 1, 0])  # Index values range from (0: number of columns) in "x"!

Now if I wanted to add a certain value c to the columns of x depending on the indices idx, I would do the following:
x[range(3), idx] += c

To get:
x = np.array([[  0,  1,  0,  c,  0],
              [  0,  c,  0,  1,  0],
              [1+c,  0,  0,  0,  0]])

But what if I wanted to add the value c to every other column index in x, rather than the exact indices in idx?
The desired outcome (based on the above example) should be:
x = np.array([[c, 1+c,  c,   0,  c],
              [c,   0,  c, 1+c,  c],
              [1,   c,  c,   c,  c]])

Output:
Create a boolean array to use as mask:
# set up default mask
m = np.ones(x.shape, dtype=bool)
# update mask
m[np.arange(m.shape[0]), idx] = False
# perform boolean indexing
x[m] += c

Output (c=9):
array([[ 9, 10,  9,  0,  9],
       [ 9,  0,  9, 10,  9],
       [ 1,  9,  9,  9,  9]])

m:
array([[ True,  True,  True, False,  True],
       [ True, False,  True,  True,  True],
       [False,  True,  True,  True,  True]])


Instruction: CNN training loss is so unstable

Input:
my CNN network
Above is my config of the network.
l am training a CNN network on picture size of 192*192.
my target is a classification network of 11 kinds.
However, the loss and the accuracy on testing dataset appears to be very unstable. l have to run 15+ epochs to get a stable accuracy and loss. The maximum accuracy is only 50%.
What can l do to improve the performance?

Output:
I would recommend you to first refer to models which are widely known like VGG-16, LeNET or VGG-19 and check out the way how the conv2D and max-pooling layers are placed.
Start with a very basic model without any batch normalization and Leaky ReLU layers. You just keep the conv2D and max pooling layers and train your model for a few epochs.
Next, try other activations like ReLU to TanH. Try Changing the max pooling to average pooling.
If you are solving a classification problem then use the softmax layer at the end. Also, introduce Dense layer(s) after flattening.
Your dataset should be large and also the target should be one-hot encoded if you wish to use the softmax layer.