facebookresearch / ConvNeXt

Code release for ConvNeXt model
MIT License
5.78k stars 696 forks source link

Number Of Classes #110

Closed SerdarHelli closed 2 years ago

SerdarHelli commented 2 years ago

I' am trying to run the below code. However, the number of classes doesnt be changed, although I write it on script

!python -m torch.distributed.launch --nproc_per_node=1 main.py \ --model convnext_small --drop_path 0.4 --input_size 224 \ --nb_classes=7 \ --batch_size 32 --lr 4e-3 --update_freq 1 \ --model_ema true --model_ema_eval true --epochs 50 \ --data_path /content/data/ \ --finetune https://dl.fbaipublicfiles.com/convnext/convnext_small_22k_224.pth \ --eval_data_path /content/data/ \ --output_dir /content/

jrsykes commented 2 years ago

HI, I'm having the same problem while trying to train a ConvNeXt tiny model by loading the model from torchvision.models. Pytorch version: 1.13.0.dev20220624

I use the following two lines of code to load the model and change the number of output nodes:

num_classes=2
model_ft = models.convnext_tiny(weights=ConvNeXt_Tiny_Weights.DEFAULT) model_ft.classifier[2].out_features = num_classes

And when I print this layer of the mode I get:

print(model_ft.classifier[2])

Linear(in_features=768, out_features=2, bias=True)

Which suggests that the change had been made. However, when I train the model, the output has dimensions of 42 x 1,000. i.e. batch_size x n classes in ImageNet:

batch_size=42 outputs = model(inputs) print(outputs.size())

torch.Size([42, 1000])

Any thoughts on how solve this problem?

Cheers, Jamie

p.s. it seems like the issue might be that the number of classes is hard coded as 1000 in: pytorch/vision/tree/main/torchvision/models/convnext.py Lines 90:100

class ConvNeXt(nn.Module): def init( self, block_setting: List[CNBlockConfig], stochastic_depth_prob: float = 0.0, layer_scale: float = 1e-6, num_classes: int = 1000, block: Optional[Callable[..., nn.Module]] = None, norm_layer: Optional[Callable[..., nn.Module]] = None, **kwargs: Any, ) -> None:

SerdarHelli commented 2 years ago

Thanks <3