Skip to content

Instantly share code, notes, and snippets.

@hasithsura

hasithsura/model.py

Created January 9, 2020 07:36
Show Gist options
  • Save hasithsura/48d92e238a6562a2dbc4cc4cea0a76f2 to your computer and use it in GitHub Desktop.
Save hasithsura/48d92e238a6562a2dbc4cc4cea0a76f2 to your computer and use it in GitHub Desktop.
Download ZIP
Raw
model.py
from torchvision.models import resnet34
import torch
import torch.nn as nn
import torch.optim as optim
if torch.cuda.is_available():
device=torch.device('cuda:0')
else:
device=torch.device('cpu')
resnet_model = resnet34(pretrained=True)
resnet_model.fc = nn.Linear(512,50)
resnet_model.conv1 = nn.Conv2d(1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
resnet_model = resnet_model.to(device)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment