<jupyter_start><jupyter_text><jupyter_code>imagenet_data = torchvision.datasets.('temp',download=True, transform=_transforms)
data_loader = torch.utils.data.DataLoader(imagenet_data,
                                          batch_size=4,
                                          shuffle=True,
                                          num_workers=4)

resnet18 = torchvision.models.resnet18(pretrained=True)<jupyter_output>Downloading: "https://download.pytorch.org/models/resnet18-5c106cde.pth" to /root/.cache/torch/checkpoints/resnet18-5c106cde.pth
100%|██████████| 44.7M/44.7M [00:01<00:00, 36.1MB/s]<jupyter_text>https://stackoverflow.com/questions/52548174/how-to-remove-the-last-fc-layer-from-a-resnet-model-in-pytorch<jupyter_code>class CustomHead(nn.Module):
    def __init__(self, n, out):
        super().__init__()
        
        self.lin = nn.Linear(n, out)
        self.act = nn.Softmax(0)
    
    def forward(x):
        x = self.lin(x)
        return self.act(x)
model = nn.Sequential(
    OrderedDict([
        ('pre-trained', nn.Sequential(*list(resnet18.children())[:-1])),
        ('head', CustomHead(512,10))
    ])
)
for i in data_loader:
    print(i)
    X, y = i
    out = model(X)
    break<jupyter_output><empty_output>
StarcoderdataJupyterStructuredCleanDedup