mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-10-25 01:40:08 +00:00
648f806315
* add example of self-supervised SimCLR training * simclr v2, replace nvidia dali dataloader * updated * sync to latest code writing style * sync to latest code writing style and modify README * detail README & standardize dataset path
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
|