mirror of https://github.com/hpcaitech/ColossalAI
20 lines
504 B
Python
20 lines
504 B
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from .Backbone import backbone
|
|
|
|
class Linear_eval(nn.Module):
|
|
|
|
def __init__(self, model='resnet18', class_num=10, **kwargs):
|
|
super().__init__()
|
|
|
|
self.backbone = backbone(model, **kwargs)
|
|
self.backbone.requires_grad_(False)
|
|
self.fc = nn.Linear(self.backbone.output_dim, class_num)
|
|
|
|
def forward(self, x):
|
|
|
|
out = self.backbone(x)
|
|
out = self.fc(out)
|
|
return out
|